zwhui 2 년 전
부모
커밋
ac451ff2ca

+ 11 - 0
midjourney/src/main/java/com/yhlxj/dao/mapper/midjourney/MidjourneyWhiteMapper.java

@@ -0,0 +1,11 @@
+package com.yhlxj.dao.mapper.midjourney;
+
+import com.baomidou.mybatisplus.core.mapper.BaseMapper;
+import com.yhlxj.dao.model.entity.MidjourneyWhite;
+
+/**
+ * @author zwhui
+ * @date 2024/7/12 11:02
+ */
+public interface MidjourneyWhiteMapper extends BaseMapper<MidjourneyWhite> {
+}

+ 13 - 0
midjourney/src/main/java/com/yhlxj/dao/model/entity/MidjourneyWhite.java

@@ -0,0 +1,13 @@
+package com.yhlxj.dao.model.entity;
+
+import com.yhlxj.dao.model.BaseEntity;
+import lombok.Data;
+
+/**
+ * @author zwhui
+ * @date 2024/7/12 11:01
+ */
+@Data
+public class MidjourneyWhite extends BaseEntity {
+    private Long userId;
+}

+ 2 - 2
midjourney/src/main/java/com/yhlxj/redis/RedisService.java

@@ -231,7 +231,7 @@ public class RedisService {
      * @param time 时间(秒)
      * @return true成功 false失败
      */
-    public Boolean hmset(String key, Map<String, Object> map, Long time) {
+    public Boolean hmset(String key, Map<Object, Object> map, Long time) {
         try {
             redisTemplate.opsForHash().putAll(key, map);
             if (time > 0) {
@@ -751,7 +751,7 @@ public class RedisService {
         MIDJOURNEY_EXPIRE_TIME("midjourney:expire:time:", "midjourney expire time", 60 * 60 * 48L),
         MIDJOURNEY_USER("midjourney:user:", "midjourney user", 60 * 60 * 2L),
         MIDJOURNEY_CONVERSATION("midjourney:conversation:", "midjourney conversation", 60 * 60 * 2L),
-        MIDJOURNEY_ACCOUNT("midjourney:account:", "midjourney account", 60 * 60 * 48L),
+        MIDJOURNEY_ACCOUNT("midjourney:account:", "midjourney account", 60 * 60 * 2L),
         MIDJOURNEY_QUERY("midjourney:query:", "midjourney query", 60 * 60 * 2L),
         MIDJOURNEY_PROGRESS("midjourney:task:progress:","task progress 任务进度", 60 * 60 * 1L),
         MIDJOURNEY_USER_TIME("midjourney:user:time:", "midjourney user time", 60 * 60 * 24L),

+ 7 - 2
midjourney/src/main/java/com/yhlxj/service/midjourney/impl/MidJourneyAccountServiceImpl.java

@@ -74,6 +74,11 @@ public class MidJourneyAccountServiceImpl extends ServiceImpl<MidjourneyAccountM
         Long instanceId = getAccount(midjourneyAccount.getId());
         midjourneyAccount.setInstanceId(instanceId);
         baseMapper.updateById(midjourneyAccount);
+        String key = RedisService.key.MIDJOURNEY_ACCOUNT.getName();
+        Map<Object, Object> accountsUsage = redisService.hmget(key);
+        if (MapUtil.isNotEmpty(accountsUsage)) {
+            redisService.hset(key,midjourneyAccount.getInstanceId().toString(),accountsUsage.values().toArray()[0]);
+        }
     }
 
     @Override
@@ -325,7 +330,7 @@ public class MidJourneyAccountServiceImpl extends ServiceImpl<MidjourneyAccountM
                             }
                         });
                     }else {
-                        redisService.hdel(RedisService.key.MIDJOURNEY_ACCOUNT.getName(), instanceId);
+                        redisService.hdel(RedisService.key.MIDJOURNEY_ACCOUNT.getName(), instanceId.toString());
                     }
 
                 }
@@ -340,7 +345,7 @@ public class MidJourneyAccountServiceImpl extends ServiceImpl<MidjourneyAccountM
         list.forEach((account)->{
             if (!ids.contains(account.getInstanceId().toString())) {
                 removeById(account.getId());
-                redisService.hdel(RedisService.key.MIDJOURNEY_ACCOUNT.getName(),account.getInstanceId());
+                redisService.hdel(RedisService.key.MIDJOURNEY_ACCOUNT.getName(),account.getInstanceId().toString());
             }
         });
 

+ 14 - 5
midjourney/src/main/java/com/yhlxj/service/midjourney/impl/MidjourneyServiceImpl.java

@@ -20,6 +20,7 @@ import com.yhlxj.dao.mapper.GroupsRelationMapper;
 import com.yhlxj.dao.mapper.midjourney.MidjourneyApiChannelMapper;
 import com.yhlxj.dao.mapper.midjourney.MidjourneyUserConversationMapper;
 import com.yhlxj.dao.mapper.midjourney.MidjourneyUserMapper;
+import com.yhlxj.dao.mapper.midjourney.MidjourneyWhiteMapper;
 import com.yhlxj.dao.model.dto.BlendDimensions;
 import com.yhlxj.dao.model.dto.MessageButton;
 import com.yhlxj.dao.model.dto.MjVersionStatus;
@@ -110,6 +111,7 @@ public class MidjourneyServiceImpl implements MidjourneyService {
 
     public static final Map<String, MjVersionStatus> MJ_VERSION_STATUS_MAP = new ConcurrentHashMap<>();
 
+    private final MidjourneyWhiteMapper midjourneyWhiteMapper;
 
     @Override
     public MidjourneyUser getUser(String userToken) {
@@ -570,7 +572,9 @@ public class MidjourneyServiceImpl implements MidjourneyService {
         if (count == null) {
             redisService.set(key, 1L,RedisService.key.MIDJOURNEY_USER_TIME.getTimeout());
         }else {
-            if (redisService.incr(key, 1L) > 200) {
+            Integer white = midjourneyWhiteMapper.selectCount(Wrappers.lambdaQuery(MidjourneyWhite.class).eq(MidjourneyWhite::getUserId, user.getId()));
+            int limit = white == 0 ? 200 : 400;
+            if (redisService.incr(key, 1L) > limit) {
                 //24小时内提交次数超过200次 拉黑
                 user.setIsBlack(true);
                 midjourneyUserMapper.updateById(user);
@@ -586,15 +590,15 @@ public class MidjourneyServiceImpl implements MidjourneyService {
     }
 
     public String getAccountWithMinUsage(Long instanceId) {
-        String key = RedisService.key.MIDJOURNEY_ACCOUNT.getName();
+        RedisService.key key = RedisService.key.MIDJOURNEY_ACCOUNT;
         try {
             // 获取所有账号ID和使用次数
-            Map<Object, Object> accountsUsage = redisService.hmget(key);
+            Map<Object, Object> accountsUsage = redisService.hmget(key.getName());
 
             if (accountsUsage == null || accountsUsage.isEmpty()) {
                 accountsUsage = midjourneyAccountService.list(Wrappers.lambdaQuery(MidjourneyAccount.class).eq(MidjourneyAccount::getStatus,Boolean.TRUE)).stream().collect(Collectors.toMap(k -> k.getInstanceId().toString(), v -> 0));
                 // 保存所有账号ID和使用次数
-                redisService.hmset(key, accountsUsage);
+                redisService.hmset(key.getName(), accountsUsage,key.getTimeout());
             }
 
             // 找到使用次数最少的账号ID
@@ -716,7 +720,7 @@ public class MidjourneyServiceImpl implements MidjourneyService {
                 //失败返还次数
                 if ("FAILURE".equals(conversation.getStatus())){
                     if (StringUtils.equals("未知频道",conversation.getFailReason())) {
-                        redisService.hdel(RedisService.key.MIDJOURNEY_ACCOUNT.getName(),dbConversation.getInstanceId());
+                        redisService.hdel(RedisService.key.MIDJOURNEY_ACCOUNT.getName(),dbConversation.getInstanceId().toString());
                         String finalAccountWithMinUsage = dbConversation.getInstanceId().toString();
                         TASK_EXECUTOR.execute(() -> {
                             MidjourneyAccount account = midjourneyAccountService.getOne(Wrappers.lambdaQuery(MidjourneyAccount.class).eq(MidjourneyAccount::getInstanceId, finalAccountWithMinUsage).last("limit 1"));
@@ -824,6 +828,11 @@ public class MidjourneyServiceImpl implements MidjourneyService {
         MidjourneyUserConversation conversation = JSONUtil.toBean(jsons, MidjourneyUserConversation.class);
         conversation.setUserId(jsons.getLong("state"));
         conversation.setTaskId(jsons.getLong("id"));
+        JSONObject jsonObject = new JSONObject(conversation.getProperties());
+        if ("SUBMITTED".equals(conversation.getStatus())){
+            jsonObject.set("messageContent","**"+conversation.getPrompt()+"** - <@mj>" + (conversation.getMode() == 1 ? "(fast)" : "(relax)"));
+            conversation.setProperties(jsonObject.toString());
+        }
         if (StringUtil.isNotBlank(conversation.getImageUrl())) {
             log.info("获取到的图片 url:{}", conversation.getImageUrl());
             ImgSize imgSize = TASK_IMG_SIZE_MAP.get(conversation.getTaskId());

+ 1 - 1
netflix-service/src/main/java/com/cyksj/service/midjourney/impl/MidJourneyAccountServiceImpl.java

@@ -80,7 +80,7 @@ public class MidJourneyAccountServiceImpl extends ServiceImpl<MidjourneyAccountM
         baseMapper.deleteById(id);
         //删除plus服务账号信息
         delAcount(midjourneyAccount.getInstanceId());
-        redisService.hdel(RedisService.key.MIDJOURNEY_ACCOUNT.getName(),midjourneyAccount.getInstanceId());
+        redisService.hdel(RedisService.key.MIDJOURNEY_ACCOUNT.getName(),midjourneyAccount.getInstanceId().toString());
     }
 
     @Override

+ 1 - 1
netflix-service/src/main/java/com/cyksj/service/midjourney/impl/MidjourneyServiceImpl.java

@@ -415,7 +415,7 @@ public class MidjourneyServiceImpl implements MidjourneyService {
                 //失败返还次数
                 if ("FAILURE".equals(conversation.getStatus())){
                     if (StringUtils.equals("未知频道",conversation.getFailReason())) {
-                        redisService.hdel(RedisService.key.MIDJOURNEY_ACCOUNT.getName(),dbConversation.getInstanceId());
+                        redisService.hdel(RedisService.key.MIDJOURNEY_ACCOUNT.getName(),dbConversation.getInstanceId().toString());
                         String finalAccountWithMinUsage = dbConversation.getInstanceId().toString();
                         TASK_EXECUTOR.execute(() -> {
                             MidjourneyAccount account = midjourneyAccountService.getOne(Wrappers.lambdaQuery(MidjourneyAccount.class).eq(MidjourneyAccount::getInstanceId, finalAccountWithMinUsage).last("limit 1"));