diff --git a/src/main/java/work/slhaf/chatai/App.java b/src/main/java/work/slhaf/chatai/App.java index fd07507..027d283 100644 --- a/src/main/java/work/slhaf/chatai/App.java +++ b/src/main/java/work/slhaf/chatai/App.java @@ -8,6 +8,7 @@ import net.mamoe.mirai.event.events.FriendMessageEvent; import net.mamoe.mirai.event.events.GroupMessageEvent; import net.mamoe.mirai.event.events.MessageEvent; import work.slhaf.chatai.chat.constant.Constant; +import work.slhaf.chatai.chat.util.ChatUtil; import work.slhaf.chatai.config.Config; import work.slhaf.chatai.config.CustomCommandTemplate; import work.slhaf.chatai.config.ModelConfigTemplate; @@ -45,6 +46,7 @@ public final class App extends JavaPlugin { try { Config config = Config.ConfigLoader.load(); OCRUtil.load(); + ChatUtil.loadCustom(); owner = config.getOwner().substring(1); bot = config.getBot().substring(1); diff --git a/src/main/java/work/slhaf/chatai/chat/ChatClient.java b/src/main/java/work/slhaf/chatai/chat/ChatClient.java index bb7bd24..3e0f405 100644 --- a/src/main/java/work/slhaf/chatai/chat/ChatClient.java +++ b/src/main/java/work/slhaf/chatai/chat/ChatClient.java @@ -12,6 +12,7 @@ import work.slhaf.chatai.chat.pojo.ChatBody; import work.slhaf.chatai.chat.pojo.ChatResponse; import work.slhaf.chatai.chat.pojo.Message; import work.slhaf.chatai.chat.pojo.PrimaryChatResponse; +import work.slhaf.chatai.chat.util.ChatUtil; import java.util.ArrayList; import java.util.List; @@ -57,13 +58,23 @@ public class ChatClient { } public void setPromotion(String promotion) { - if (!promotion.equals("null")) { - Message message = Message.builder() - .role(Constant.Character.SYSTEM) - .content(promotion) - .build(); - messages.add(0, message); + if (promotion.equals("null")) { + return; } + if (promotion.startsWith(Constant.Model.PREFIX_PROMOTION)) { + String customName = promotion.substring(Constant.Model.PREFIX_PROMOTION.length()); + ArrayList list = ChatUtil.customMessages.get(customName); + if (list == null) { + return; + } + messages.addAll(list); + return; + } + Message message = Message.builder() + .role(Constant.Character.SYSTEM) + .content(promotion) + .build(); + messages.add(0, message); } public ChatResponse runChat(String content) { diff --git a/src/main/java/work/slhaf/chatai/chat/constant/Constant.java b/src/main/java/work/slhaf/chatai/chat/constant/Constant.java index b6bfe30..83b17dd 100644 --- a/src/main/java/work/slhaf/chatai/chat/constant/Constant.java +++ b/src/main/java/work/slhaf/chatai/chat/constant/Constant.java @@ -12,6 +12,7 @@ public class Constant { public static final String GLM_4_FLASH = "glm-4-flash"; public static final String GLM_4_0520 = "glm-4-0520"; public static final String GLM_4_PLUS = "glm-4-Plus"; + public static final String PREFIX_PROMOTION = "custom::"; } public static class Status{ @@ -62,5 +63,6 @@ public class Constant { public static final String CONFIG_ALIOCR_PATH = "./config/ChatAI/aliocr.yaml"; public static final String CONFIG_MODEL_PATH = "./config/ChatAI/model.yaml"; public static final String CONFIG_CUSTOM_PATH = "./config/ChatAI/custom.yaml"; + public static final String DATA_CUSTOM_PATH = "./data/ChatAI/custom/"; } } diff --git a/src/main/java/work/slhaf/chatai/chat/util/ChatUtil.java b/src/main/java/work/slhaf/chatai/chat/util/ChatUtil.java index 35ffd2d..4751ad7 100644 --- a/src/main/java/work/slhaf/chatai/chat/util/ChatUtil.java +++ b/src/main/java/work/slhaf/chatai/chat/util/ChatUtil.java @@ -1,13 +1,23 @@ package work.slhaf.chatai.chat.util; +import com.alibaba.fastjson.JSON; +import com.alibaba.fastjson.JSONArray; +import com.alibaba.fastjson.JSONObject; import lombok.extern.slf4j.Slf4j; import work.slhaf.chatai.chat.ChatClient; import work.slhaf.chatai.chat.constant.Constant; import work.slhaf.chatai.chat.pojo.ChatResponse; +import work.slhaf.chatai.chat.pojo.Message; import work.slhaf.chatai.config.Config; import work.slhaf.chatai.config.CustomCommandTemplate; import work.slhaf.chatai.config.ModelConfigTemplate; +import java.io.BufferedReader; +import java.io.File; +import java.io.FileReader; +import java.io.IOException; +import java.util.ArrayList; +import java.util.Arrays; import java.util.HashMap; import java.util.Objects; @@ -16,15 +26,53 @@ public class ChatUtil { private static final HashMap chatClients = new HashMap<>(); private static final HashMap latestTime = new HashMap<>(); + public static final HashMap> customMessages = new HashMap<>(); private ChatUtil() { - } static { launchCleanerThread(); } + public static void loadCustom() throws IOException { + File filePath = new File(Constant.Path.DATA_CUSTOM_PATH); + File[] files = filePath.listFiles(); + if (files != null) { + for (File file : files) { + StringBuilder content = new StringBuilder(); + //read files + BufferedReader br = new BufferedReader(new FileReader(file)); + String line; + while ((line = br.readLine()) != null) { + content.append(line); + } + //save data + try { + JSONObject jsonObject = JSON.parseObject(content.toString()); + String name = jsonObject.getString("name"); + JSONArray messages = jsonObject.getJSONArray("messages"); + ArrayList list = new ArrayList<>(); + for (int i = 0; i < messages.size(); i++) { + Message message = new Message(); + JSONObject messageObject = messages.getJSONObject(i); + message.setRole(messageObject.getString("role")); + message.setContent(messageObject.getString("content")); + list.add(message); + } + customMessages.put(name, list); + } catch (Exception e) { + log.error("读取出错,请检查预设文件"); + log.error(e.getMessage()); + } + br.close(); + log.info("预设文件加载完毕: {}", Arrays.toString(customMessages.keySet().toArray(new String[0]))); + } + } else { + log.info("未找到预设文件"); + } + } + private static void launchCleanerThread() { Config config = Config.ConfigLoader.getConfig(); long timeout = Long.parseLong(config.getTimeout().substring(1)); @@ -53,7 +101,7 @@ public class ChatUtil { } public static String chat(String id, String content, String command) { - String chatId = id + "-" + command + "-" + System.currentTimeMillis(); + String chatId = id + "-" + command; ChatResponse chatResponse = null; synchronized (chatClients) { if (chatClients.containsKey(chatId)) { @@ -83,7 +131,7 @@ public class ChatUtil { String url = modelConfigTemplate.getBase_url(); String apikey = modelConfigTemplate.getApikey(); - ChatClient chatClient = new ChatClient(chatId, url, apikey, customModel,top_p,temperature,max_tokens); + ChatClient chatClient = new ChatClient(chatId, url, apikey, customModel, top_p, temperature, max_tokens); chatClient.setPromotion(customPromotion); chatClients.put(chatId, chatClient); log.info("final content: {}", content); diff --git a/src/main/java/work/slhaf/chatai/config/Config.java b/src/main/java/work/slhaf/chatai/config/Config.java index 8785ca2..6c148de 100644 --- a/src/main/java/work/slhaf/chatai/config/Config.java +++ b/src/main/java/work/slhaf/chatai/config/Config.java @@ -78,6 +78,11 @@ public class Config { } public static Config load() throws IOException { + //create custom data files + File customDirs = new File(Constant.Path.DATA_CUSTOM_PATH); + customDirs.mkdirs(); + + //create config files config = new Config(); LoaderOptions loaderOptions = new LoaderOptions(); TagInspector tagInspector = tag -> true; diff --git a/src/test/java/MyTest.java b/src/test/java/MyTest.java index c28abf7..9f5623b 100644 --- a/src/test/java/MyTest.java +++ b/src/test/java/MyTest.java @@ -14,6 +14,8 @@ import org.apache.http.message.BasicHeader; import org.apache.http.util.EntityUtils; import org.junit.Test; import work.slhaf.chatai.App; +import work.slhaf.chatai.chat.util.ChatUtil; +import work.slhaf.chatai.config.Config; import java.io.IOException; import java.util.Scanner; @@ -121,7 +123,7 @@ public class MyTest { @Test - public void terminalTest() throws InterruptedException { + public void terminalTest() { MiraiConsoleImplementationTerminal terminal = new MiraiConsoleImplementationTerminal(); MiraiConsoleTerminalLoader.INSTANCE.startAsDaemon(terminal); PluginManager.INSTANCE.loadPlugin(App.INSTANCE); @@ -134,4 +136,10 @@ public class MyTest { } } } + + @Test + public void customTest() throws IOException { + Config.ConfigLoader.load(); + ChatUtil.loadCustom(); + } }