From dea23f13ef2d4918080ce8aeee1ee908188cdb19 Mon Sep 17 00:00:00 2001
From: ageer <ageerle@163.com>
Date: 星期一, 01 四月 2024 22:21:29 +0800
Subject: [PATCH] v1.0.3
---
ruoyi-modules/ruoyi-midjourney/src/main/java/com/xmzs/midjourney/controller/SubmitController.java | 477 ++++++++++++++++++++++++++++++++++++-----------------------
1 files changed, 292 insertions(+), 185 deletions(-)
diff --git a/ruoyi-modules/ruoyi-midjourney/src/main/java/com/xmzs/midjourney/controller/SubmitController.java b/ruoyi-modules/ruoyi-midjourney/src/main/java/com/xmzs/midjourney/controller/SubmitController.java
index 3973087..fb3fe29 100644
--- a/ruoyi-modules/ruoyi-midjourney/src/main/java/com/xmzs/midjourney/controller/SubmitController.java
+++ b/ruoyi-modules/ruoyi-midjourney/src/main/java/com/xmzs/midjourney/controller/SubmitController.java
@@ -2,15 +2,15 @@
import cn.hutool.core.text.CharSequenceUtil;
import cn.hutool.core.util.RandomUtil;
+import cn.hutool.json.JSONUtil;
+import com.xmzs.common.chat.constant.OpenAIConst;
+import com.xmzs.common.core.domain.model.LoginUser;
+import com.xmzs.common.core.exception.base.BaseException;
+import com.xmzs.common.satoken.utils.LoginHelper;
import com.xmzs.midjourney.Constants;
import com.xmzs.midjourney.ProxyProperties;
import com.xmzs.midjourney.ReturnCode;
-import com.xmzs.midjourney.dto.BaseSubmitDTO;
-import com.xmzs.midjourney.dto.SubmitBlendDTO;
-import com.xmzs.midjourney.dto.SubmitChangeDTO;
-import com.xmzs.midjourney.dto.SubmitDescribeDTO;
-import com.xmzs.midjourney.dto.SubmitImagineDTO;
-import com.xmzs.midjourney.dto.SubmitSimpleChangeDTO;
+import com.xmzs.midjourney.dto.*;
import com.xmzs.midjourney.enums.TaskAction;
import com.xmzs.midjourney.enums.TaskStatus;
import com.xmzs.midjourney.enums.TranslateWay;
@@ -26,215 +26,322 @@
import com.xmzs.midjourney.util.MimeTypeUtils;
import com.xmzs.midjourney.util.SnowFlake;
import com.xmzs.midjourney.util.TaskChangeParams;
+import com.xmzs.system.domain.bo.ChatMessageBo;
+import com.xmzs.system.service.IChatMessageService;
+import com.xmzs.system.service.IChatService;
+import com.xmzs.system.service.ISseService;
import eu.maxschuster.dataurl.DataUrl;
import eu.maxschuster.dataurl.DataUrlSerializer;
import eu.maxschuster.dataurl.IDataUrlSerializer;
import io.swagger.annotations.Api;
import io.swagger.annotations.ApiOperation;
import lombok.RequiredArgsConstructor;
+import lombok.extern.slf4j.Slf4j;
+import okhttp3.OkHttpClient;
+import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.beans.factory.annotation.Value;
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 java.io.IOException;
import java.net.MalformedURLException;
import java.util.ArrayList;
import java.util.List;
import java.util.Optional;
import java.util.Set;
+import java.util.concurrent.TimeUnit;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
+
+import okhttp3.*;
@Api(tags = "浠诲姟鎻愪氦")
@RestController
@RequestMapping("/mj/submit")
@RequiredArgsConstructor
+@Slf4j
public class SubmitController {
- private final TranslateService translateService;
- private final TaskStoreService taskStoreService;
- private final ProxyProperties properties;
- private final TaskService taskService;
+ private final TranslateService translateService;
+ private final ProxyProperties properties;
+ private final TaskService taskService;
+ private final TaskStoreService taskStoreService;
+ @Value("${chat.apiKey}")
+ private String apiKey;
+ @Value("${chat.apiHost}")
+ private String apiHost;
+ @Autowired
+ private IChatService chatService;
+ @Autowired
+ private IChatMessageService chatMessageService;
+ @Autowired
+ private ISseService sseService;
- @ApiOperation(value = "鎻愪氦Imagine浠诲姟")
- @PostMapping("/imagine")
- public SubmitResultVO imagine(@RequestBody SubmitImagineDTO imagineDTO) {
- String prompt = imagineDTO.getPrompt();
- if (CharSequenceUtil.isBlank(prompt)) {
- return SubmitResultVO.fail(ReturnCode.VALIDATION_ERROR, "prompt涓嶈兘涓虹┖");
- }
- prompt = prompt.trim();
- Task task = newTask(imagineDTO);
- task.setAction(TaskAction.IMAGINE);
- task.setPrompt(prompt);
- String promptEn = translatePrompt(prompt);
- try {
- BannedPromptUtils.checkBanned(promptEn);
- } catch (BannedPromptException e) {
- return SubmitResultVO.fail(ReturnCode.BANNED_PROMPT, "鍙兘鍖呭惈鏁忔劅璇�")
- .setProperty("promptEn", promptEn).setProperty("bannedWord", e.getMessage());
- }
- List<String> base64Array = Optional.ofNullable(imagineDTO.getBase64Array()).orElse(new ArrayList<>());
- if (CharSequenceUtil.isNotBlank(imagineDTO.getBase64())) {
- base64Array.add(imagineDTO.getBase64());
- }
- List<DataUrl> dataUrls;
- try {
- dataUrls = ConvertUtils.convertBase64Array(base64Array);
- } catch (MalformedURLException e) {
- return SubmitResultVO.fail(ReturnCode.VALIDATION_ERROR, "base64鏍煎紡閿欒");
- }
- task.setPromptEn(promptEn);
- task.setDescription("/imagine " + prompt);
- return this.taskService.submitImagine(task, dataUrls);
- }
+ @ApiOperation(value = "鎻愪氦Imagine浠诲姟")
+ @PostMapping("/imagine")
+ public SubmitResultVO imagine(@RequestBody SubmitImagineDTO imagineDTO) {
+ String prompt = imagineDTO.getPrompt();
+ if (CharSequenceUtil.isBlank(prompt)) {
+ return SubmitResultVO.fail(ReturnCode.VALIDATION_ERROR, "prompt涓嶈兘涓虹┖");
+ }
+ prompt = prompt.trim();
+ Task task = newTask(imagineDTO);
+ task.setAction(TaskAction.IMAGINE);
+ task.setPrompt(prompt);
+ String promptEn = translatePrompt(prompt);
+ try {
+ BannedPromptUtils.checkBanned(promptEn);
+ } catch (BannedPromptException e) {
+ return SubmitResultVO.fail(ReturnCode.BANNED_PROMPT, "鍙兘鍖呭惈鏁忔劅璇�")
+ .setProperty("promptEn", promptEn).setProperty("bannedWord", e.getMessage());
+ }
+ List<String> base64Array = Optional.ofNullable(imagineDTO.getBase64Array()).orElse(new ArrayList<>());
+ if (CharSequenceUtil.isNotBlank(imagineDTO.getBase64())) {
+ base64Array.add(imagineDTO.getBase64());
+ }
+ List<DataUrl> dataUrls;
+ try {
+ dataUrls = ConvertUtils.convertBase64Array(base64Array);
+ } catch (MalformedURLException e) {
+ return SubmitResultVO.fail(ReturnCode.VALIDATION_ERROR, "base64鏍煎紡閿欒");
+ }
+ task.setPromptEn(promptEn);
+ task.setDescription("/imagine " + prompt);
+ return this.taskService.submitImagine(task, dataUrls);
+ }
- @ApiOperation(value = "缁樺浘鍙樺寲-simple")
- @PostMapping("/simple-change")
- public SubmitResultVO simpleChange(@RequestBody SubmitSimpleChangeDTO simpleChangeDTO) {
- TaskChangeParams changeParams = ConvertUtils.convertChangeParams(simpleChangeDTO.getContent());
- if (changeParams == null) {
- return SubmitResultVO.fail(ReturnCode.VALIDATION_ERROR, "content鍙傛暟閿欒");
- }
- SubmitChangeDTO changeDTO = new SubmitChangeDTO();
- changeDTO.setAction(changeParams.getAction());
- changeDTO.setTaskId(changeParams.getId());
- changeDTO.setIndex(changeParams.getIndex());
- changeDTO.setState(simpleChangeDTO.getState());
- changeDTO.setNotifyHook(simpleChangeDTO.getNotifyHook());
- return change(changeDTO);
- }
+ @ApiOperation(value = "缁樺浘鍙樺寲-simple")
+ @PostMapping("/simple-change")
+ public SubmitResultVO simpleChange(@RequestBody SubmitSimpleChangeDTO simpleChangeDTO) {
+ TaskChangeParams changeParams = ConvertUtils.convertChangeParams(simpleChangeDTO.getContent());
+ if (changeParams == null) {
+ return SubmitResultVO.fail(ReturnCode.VALIDATION_ERROR, "content鍙傛暟閿欒");
+ }
+ SubmitChangeDTO changeDTO = new SubmitChangeDTO();
+ changeDTO.setAction(changeParams.getAction());
+ changeDTO.setTaskId(changeParams.getId());
+ changeDTO.setIndex(changeParams.getIndex());
+ changeDTO.setState(simpleChangeDTO.getState());
+ changeDTO.setNotifyHook(simpleChangeDTO.getNotifyHook());
+ return change(changeDTO);
+ }
- @ApiOperation(value = "缁樺浘鍙樺寲")
- @PostMapping("/change")
- public SubmitResultVO change(@RequestBody SubmitChangeDTO changeDTO) {
- if (CharSequenceUtil.isBlank(changeDTO.getTaskId())) {
- return SubmitResultVO.fail(ReturnCode.VALIDATION_ERROR, "taskId涓嶈兘涓虹┖");
- }
- if (!Set.of(TaskAction.UPSCALE, TaskAction.VARIATION, TaskAction.REROLL).contains(changeDTO.getAction())) {
- return SubmitResultVO.fail(ReturnCode.VALIDATION_ERROR, "action鍙傛暟閿欒");
- }
- String description = "/up " + changeDTO.getTaskId();
- if (TaskAction.REROLL.equals(changeDTO.getAction())) {
- description += " R";
- } else {
- description += " " + changeDTO.getAction().name().charAt(0) + changeDTO.getIndex();
- }
- if (TaskAction.UPSCALE.equals(changeDTO.getAction())) {
- TaskCondition condition = new TaskCondition().setDescription(description);
- Task existTask = this.taskStoreService.findOne(condition);
- if (existTask != null) {
- return SubmitResultVO.of(ReturnCode.EXISTED, "浠诲姟宸插瓨鍦�", existTask.getId())
- .setProperty("status", existTask.getStatus())
- .setProperty("imageUrl", existTask.getImageUrl());
- }
- }
- Task targetTask = this.taskStoreService.get(changeDTO.getTaskId());
- if (targetTask == null) {
- return SubmitResultVO.fail(ReturnCode.NOT_FOUND, "鍏宠仈浠诲姟涓嶅瓨鍦ㄦ垨宸插け鏁�");
- }
- if (!TaskStatus.SUCCESS.equals(targetTask.getStatus())) {
- return SubmitResultVO.fail(ReturnCode.VALIDATION_ERROR, "鍏宠仈浠诲姟鐘舵�侀敊璇�");
- }
- if (!Set.of(TaskAction.IMAGINE, TaskAction.VARIATION, TaskAction.REROLL, TaskAction.BLEND).contains(targetTask.getAction())) {
- return SubmitResultVO.fail(ReturnCode.VALIDATION_ERROR, "鍏宠仈浠诲姟涓嶅厑璁告墽琛屽彉鍖�");
- }
- Task task = newTask(changeDTO);
- task.setAction(changeDTO.getAction());
- task.setPrompt(targetTask.getPrompt());
- task.setPromptEn(targetTask.getPromptEn());
- task.setProperty(Constants.TASK_PROPERTY_FINAL_PROMPT, targetTask.getProperty(Constants.TASK_PROPERTY_FINAL_PROMPT));
- task.setProperty(Constants.TASK_PROPERTY_PROGRESS_MESSAGE_ID, targetTask.getProperty(Constants.TASK_PROPERTY_MESSAGE_ID));
- task.setProperty(Constants.TASK_PROPERTY_DISCORD_INSTANCE_ID, targetTask.getProperty(Constants.TASK_PROPERTY_DISCORD_INSTANCE_ID));
- task.setDescription(description);
- int messageFlags = targetTask.getPropertyGeneric(Constants.TASK_PROPERTY_FLAGS);
- String messageId = targetTask.getPropertyGeneric(Constants.TASK_PROPERTY_MESSAGE_ID);
- String messageHash = targetTask.getPropertyGeneric(Constants.TASK_PROPERTY_MESSAGE_HASH);
- if (TaskAction.UPSCALE.equals(changeDTO.getAction())) {
- return this.taskService.submitUpscale(task, messageId, messageHash, changeDTO.getIndex(), messageFlags);
- } else if (TaskAction.VARIATION.equals(changeDTO.getAction())) {
- return this.taskService.submitVariation(task, messageId, messageHash, changeDTO.getIndex(), messageFlags);
- } else {
- return this.taskService.submitReroll(task, messageId, messageHash, messageFlags);
- }
- }
+ @ApiOperation(value = "缁樺浘鍙樺寲")
+ @PostMapping("/change")
+ public SubmitResultVO change(@RequestBody SubmitChangeDTO changeDTO) {
+ if (CharSequenceUtil.isBlank(changeDTO.getTaskId())) {
+ return SubmitResultVO.fail(ReturnCode.VALIDATION_ERROR, "taskId涓嶈兘涓虹┖");
+ }
+ if (!Set.of(TaskAction.UPSCALE, TaskAction.VARIATION, TaskAction.REROLL).contains(changeDTO.getAction())) {
+ return SubmitResultVO.fail(ReturnCode.VALIDATION_ERROR, "action鍙傛暟閿欒");
+ }
+ String description = "/up " + changeDTO.getTaskId();
+ if (TaskAction.REROLL.equals(changeDTO.getAction())) {
+ description += " R";
+ } else {
+ description += " " + changeDTO.getAction().name().charAt(0) + changeDTO.getIndex();
+ }
+ if (TaskAction.UPSCALE.equals(changeDTO.getAction())) {
+ TaskCondition condition = new TaskCondition().setDescription(description);
+ Task existTask = this.taskStoreService.findOne(condition);
+ if (existTask != null) {
+ return SubmitResultVO.of(ReturnCode.EXISTED, "浠诲姟宸插瓨鍦�", existTask.getId())
+ .setProperty("status", existTask.getStatus())
+ .setProperty("imageUrl", existTask.getImageUrl());
+ }
+ }
+ Task targetTask = this.taskStoreService.get(changeDTO.getTaskId());
+ if (targetTask == null) {
+ return SubmitResultVO.fail(ReturnCode.NOT_FOUND, "鍏宠仈浠诲姟涓嶅瓨鍦ㄦ垨宸插け鏁�");
+ }
+ if (!TaskStatus.SUCCESS.equals(targetTask.getStatus())) {
+ return SubmitResultVO.fail(ReturnCode.VALIDATION_ERROR, "鍏宠仈浠诲姟鐘舵�侀敊璇�");
+ }
+ if (!Set.of(TaskAction.IMAGINE, TaskAction.VARIATION, TaskAction.REROLL, TaskAction.BLEND).contains(targetTask.getAction())) {
+ return SubmitResultVO.fail(ReturnCode.VALIDATION_ERROR, "鍏宠仈浠诲姟涓嶅厑璁告墽琛屽彉鍖�");
+ }
+ Task task = newTask(changeDTO);
+ task.setAction(changeDTO.getAction());
+ task.setPrompt(targetTask.getPrompt());
+ task.setPromptEn(targetTask.getPromptEn());
+ task.setProperty(Constants.TASK_PROPERTY_FINAL_PROMPT, targetTask.getProperty(Constants.TASK_PROPERTY_FINAL_PROMPT));
+ task.setProperty(Constants.TASK_PROPERTY_PROGRESS_MESSAGE_ID, targetTask.getProperty(Constants.TASK_PROPERTY_MESSAGE_ID));
+ task.setProperty(Constants.TASK_PROPERTY_DISCORD_INSTANCE_ID, targetTask.getProperty(Constants.TASK_PROPERTY_DISCORD_INSTANCE_ID));
+ task.setDescription(description);
+ int messageFlags = targetTask.getPropertyGeneric(Constants.TASK_PROPERTY_FLAGS);
+ String messageId = targetTask.getPropertyGeneric(Constants.TASK_PROPERTY_MESSAGE_ID);
+ String messageHash = targetTask.getPropertyGeneric(Constants.TASK_PROPERTY_MESSAGE_HASH);
+ if (TaskAction.UPSCALE.equals(changeDTO.getAction())) {
+ return this.taskService.submitUpscale(task, messageId, messageHash, changeDTO.getIndex(), messageFlags);
+ } else if (TaskAction.VARIATION.equals(changeDTO.getAction())) {
+ return this.taskService.submitVariation(task, messageId, messageHash, changeDTO.getIndex(), messageFlags);
+ } else {
+ return this.taskService.submitReroll(task, messageId, messageHash, messageFlags);
+ }
+ }
- @ApiOperation(value = "鎻愪氦Describe浠诲姟")
- @PostMapping("/describe")
- public SubmitResultVO describe(@RequestBody SubmitDescribeDTO describeDTO) {
- if (CharSequenceUtil.isBlank(describeDTO.getBase64())) {
- return SubmitResultVO.fail(ReturnCode.VALIDATION_ERROR, "base64涓嶈兘涓虹┖");
- }
- IDataUrlSerializer serializer = new DataUrlSerializer();
- DataUrl dataUrl;
- try {
- dataUrl = serializer.unserialize(describeDTO.getBase64());
- } catch (MalformedURLException e) {
- return SubmitResultVO.fail(ReturnCode.VALIDATION_ERROR, "base64鏍煎紡閿欒");
- }
- Task task = newTask(describeDTO);
- task.setAction(TaskAction.DESCRIBE);
- String taskFileName = task.getId() + "." + MimeTypeUtils.guessFileSuffix(dataUrl.getMimeType());
- task.setDescription("/describe " + taskFileName);
- return this.taskService.submitDescribe(task, dataUrl);
- }
+ @ApiOperation(value = "缁樺浘鍙樺寲")
+ @PostMapping("/action")
+ public String action(@RequestBody SubmitActionDTO changeDTO) {
+ // 鏌ヨ鏄惁鏄粯璐圭敤鎴�
+ sseService.checkUserGrade();
+ // 鎵ｉ櫎鎺ュ彛璐圭敤
+ if ("upsample".equals(getAction(changeDTO.getCustomId()))) {
+ mjTaskDeduct("鏀惧ぇ", OpenAIConst.MJ_COST_TYPE2);
+ } else {
+ // Inpaint: 灞�閮ㄩ噸缁�
+ // reroll 閲嶇粯
+ // upsample 鏀惧ぇ
+ // zoom 鍙樼劍
+ // upscale 楂樻竻鏀惧ぇ
+ // variation 鍙樺寲
+ if (!"Inpaint".equals(getAction(changeDTO.getCustomId()))) {
+ mjTaskDeduct("鍙樺寲", OpenAIConst.MJ_COST_TYPE1);
+ }
+ }
+ OkHttpClient client = new OkHttpClient.Builder()
+ .connectTimeout(30, TimeUnit.SECONDS) // 杩炴帴瓒呮椂鏃堕棿
+ .writeTimeout(30, TimeUnit.SECONDS) // 鍐欏叆瓒呮椂鏃堕棿
+ .readTimeout(30, TimeUnit.SECONDS) // 璇诲彇瓒呮椂鏃堕棿
+ .build();
- @ApiOperation(value = "鎻愪氦Blend浠诲姟")
- @PostMapping("/blend")
- public SubmitResultVO blend(@RequestBody SubmitBlendDTO blendDTO) {
- List<String> base64Array = blendDTO.getBase64Array();
- if (base64Array == null || base64Array.size() < 2 || base64Array.size() > 5) {
- return SubmitResultVO.fail(ReturnCode.VALIDATION_ERROR, "base64List鍙傛暟閿欒");
- }
- if (blendDTO.getDimensions() == null) {
- return SubmitResultVO.fail(ReturnCode.VALIDATION_ERROR, "dimensions鍙傛暟閿欒");
- }
- IDataUrlSerializer serializer = new DataUrlSerializer();
- List<DataUrl> dataUrlList = new ArrayList<>();
- try {
- for (String base64 : base64Array) {
- DataUrl dataUrl = serializer.unserialize(base64);
- dataUrlList.add(dataUrl);
- }
- } catch (MalformedURLException e) {
- return SubmitResultVO.fail(ReturnCode.VALIDATION_ERROR, "base64鏍煎紡閿欒");
- }
- Task task = newTask(blendDTO);
- task.setAction(TaskAction.BLEND);
- task.setDescription("/blend " + task.getId() + " " + dataUrlList.size());
- return this.taskService.submitBlend(task, dataUrlList, blendDTO.getDimensions());
- }
+ String jsonStr = JSONUtil.toJsonStr(changeDTO);
- private Task newTask(BaseSubmitDTO base) {
- Task task = new Task();
- task.setId(System.currentTimeMillis() + RandomUtil.randomNumbers(3));
- task.setSubmitTime(System.currentTimeMillis());
- task.setState(base.getState());
- String notifyHook = CharSequenceUtil.isBlank(base.getNotifyHook()) ? this.properties.getNotifyHook() : base.getNotifyHook();
- task.setProperty(Constants.TASK_PROPERTY_NOTIFY_HOOK, notifyHook);
- task.setProperty(Constants.TASK_PROPERTY_NONCE, SnowFlake.INSTANCE.nextId());
- return task;
- }
+ MediaType mediaType = MediaType.parse("application/json");
+ okhttp3.RequestBody body = okhttp3.RequestBody.create(jsonStr, mediaType);
+ Request request = new Request.Builder()
+ .url(apiHost + "mj/submit/action")
+ .method("POST", body)
+ .header("mj-api-secret", apiKey) // 璁剧疆Authorization header
+ .build();
+ try {
+ Response response = client.newCall(request).execute();
+ return response.body().string();
+ } catch (IOException e) {
+ log.error("缁樺浘鍙樺寲澶辫触:{}", e.getMessage());
+ }
+ return null;
+ }
- private String translatePrompt(String prompt) {
- if (TranslateWay.NULL.equals(this.properties.getTranslateWay()) || CharSequenceUtil.isBlank(prompt)) {
- return prompt;
- }
- List<String> imageUrls = new ArrayList<>();
- Matcher imageMatcher = Pattern.compile("https?://[a-z0-9-_:@&?=+,.!/~*'%$]+\\x20+", Pattern.CASE_INSENSITIVE).matcher(prompt);
- while (imageMatcher.find()) {
- imageUrls.add(imageMatcher.group(0));
- }
- String paramStr = "";
- Matcher paramMatcher = Pattern.compile("\\x20+-{1,2}[a-z]+.*$", Pattern.CASE_INSENSITIVE).matcher(prompt);
- if (paramMatcher.find()) {
- paramStr = paramMatcher.group(0);
- }
- String imageStr = CharSequenceUtil.join("", imageUrls);
- String text = prompt.substring(imageStr.length(), prompt.length() - paramStr.length());
- if (CharSequenceUtil.isNotBlank(text)) {
- text = this.translateService.translateToEnglish(text).trim();
- }
- return imageStr + text + paramStr;
- }
+ public String getAction(String customId) {
+ // 妫�鏌� customId 鏄惁涓虹┖
+ if(customId == null || customId.isEmpty()) {
+ return null;
+ }
+ // 浣跨敤 "::" 鍒嗗壊瀛楃涓�
+ String[] parts = customId.split("::");
+ // "MJ", "Inpaint", "1", "4fca7c14-181c-4...", "SOLO"
+ if(customId.endsWith("SOLO")) {
+ return parts[1];
+ }
+ // 杩斿洖 "upsample" 鍊硷紝鍋囪瀹冩�绘槸鍦ㄧ涓変釜浣嶇疆
+ return parts[2];
+ }
+
+ public void mjTaskDeduct(String prompt, double cost) {
+ //鎵ｉ櫎璐圭敤
+ chatService.deductUserBalance(getUserId(), cost);
+ // 淇濆瓨娑堟伅璁板綍
+ ChatMessageBo chatMessageBo = new ChatMessageBo();
+ chatMessageBo.setUserId(getUserId());
+ chatMessageBo.setModelName("mj");
+ chatMessageBo.setContent(prompt);
+ chatMessageBo.setDeductCost(cost);
+ chatMessageBo.setTotalTokens(0);
+ chatMessageService.insertByBo(chatMessageBo);
+ }
+
+ /**
+ * 鑾峰彇鐢ㄦ埛Id
+ *
+ * @return
+ */
+ public Long getUserId() {
+ LoginUser loginUser = LoginHelper.getLoginUser();
+ if (loginUser == null) {
+ throw new BaseException("鐢ㄦ埛鏈櫥褰曪紒");
+ }
+ return loginUser.getUserId();
+ }
+
+ @ApiOperation(value = "鎻愪氦Describe浠诲姟")
+ @PostMapping("/describe")
+ public SubmitResultVO describe(@RequestBody SubmitDescribeDTO describeDTO) {
+ if (CharSequenceUtil.isBlank(describeDTO.getBase64())) {
+ return SubmitResultVO.fail(ReturnCode.VALIDATION_ERROR, "base64涓嶈兘涓虹┖");
+ }
+ IDataUrlSerializer serializer = new DataUrlSerializer();
+ DataUrl dataUrl;
+ try {
+ dataUrl = serializer.unserialize(describeDTO.getBase64());
+ } catch (MalformedURLException e) {
+ return SubmitResultVO.fail(ReturnCode.VALIDATION_ERROR, "base64鏍煎紡閿欒");
+ }
+ Task task = newTask(describeDTO);
+ task.setAction(TaskAction.DESCRIBE);
+ String taskFileName = task.getId() + "." + MimeTypeUtils.guessFileSuffix(dataUrl.getMimeType());
+ task.setDescription("/describe " + taskFileName);
+ return this.taskService.submitDescribe(task, dataUrl);
+ }
+
+ @ApiOperation(value = "鎻愪氦Blend浠诲姟")
+ @PostMapping("/blend")
+ public SubmitResultVO blend(@RequestBody SubmitBlendDTO blendDTO) {
+ List<String> base64Array = blendDTO.getBase64Array();
+ if (base64Array == null || base64Array.size() < 2 || base64Array.size() > 5) {
+ return SubmitResultVO.fail(ReturnCode.VALIDATION_ERROR, "base64List鍙傛暟閿欒");
+ }
+ if (blendDTO.getDimensions() == null) {
+ return SubmitResultVO.fail(ReturnCode.VALIDATION_ERROR, "dimensions鍙傛暟閿欒");
+ }
+ IDataUrlSerializer serializer = new DataUrlSerializer();
+ List<DataUrl> dataUrlList = new ArrayList<>();
+ try {
+ for (String base64 : base64Array) {
+ DataUrl dataUrl = serializer.unserialize(base64);
+ dataUrlList.add(dataUrl);
+ }
+ } catch (MalformedURLException e) {
+ return SubmitResultVO.fail(ReturnCode.VALIDATION_ERROR, "base64鏍煎紡閿欒");
+ }
+ Task task = newTask(blendDTO);
+ task.setAction(TaskAction.BLEND);
+ task.setDescription("/blend " + task.getId() + " " + dataUrlList.size());
+ return this.taskService.submitBlend(task, dataUrlList, blendDTO.getDimensions());
+ }
+
+ private Task newTask(BaseSubmitDTO base) {
+ Task task = new Task();
+ task.setId(System.currentTimeMillis() + RandomUtil.randomNumbers(3));
+ task.setSubmitTime(System.currentTimeMillis());
+ task.setState(base.getState());
+ String notifyHook = CharSequenceUtil.isBlank(base.getNotifyHook()) ? this.properties.getNotifyHook() : base.getNotifyHook();
+ task.setProperty(Constants.TASK_PROPERTY_NOTIFY_HOOK, notifyHook);
+ task.setProperty(Constants.TASK_PROPERTY_NONCE, SnowFlake.INSTANCE.nextId());
+ return task;
+ }
+
+ private String translatePrompt(String prompt) {
+ if (TranslateWay.NULL.equals(this.properties.getTranslateWay()) || CharSequenceUtil.isBlank(prompt)) {
+ return prompt;
+ }
+ List<String> imageUrls = new ArrayList<>();
+ Matcher imageMatcher = Pattern.compile("https?://[a-z0-9-_:@&?=+,.!/~*'%$]+\\x20+", Pattern.CASE_INSENSITIVE).matcher(prompt);
+ while (imageMatcher.find()) {
+ imageUrls.add(imageMatcher.group(0));
+ }
+ String paramStr = "";
+ Matcher paramMatcher = Pattern.compile("\\x20+-{1,2}[a-z]+.*$", Pattern.CASE_INSENSITIVE).matcher(prompt);
+ if (paramMatcher.find()) {
+ paramStr = paramMatcher.group(0);
+ }
+ String imageStr = CharSequenceUtil.join("", imageUrls);
+ String text = prompt.substring(imageStr.length(), prompt.length() - paramStr.length());
+ if (CharSequenceUtil.isNotBlank(text)) {
+ text = this.translateService.translateToEnglish(text).trim();
+ }
+ return imageStr + text + paramStr;
+ }
}
--
Gitblit v1.9.3