Update Bridge Version to 3 and Add Interrogation Feature

Author: LogicismDev

config.yml

@@ -1,16 +1,25 @@
 # The Horde Cluster URL to use to grab populated prompts from and to submit generations to.
 # This shouldn't be changed at all, unless if you run a private Horde Cluster and you know what you are doing.
-clusterURL: "https://stablehorde.net"
+clusterURL: "https://aihorde.net"
+
+# The Backup Horde Cluster URL to use if the main Horde Cluster is down
+# This shouldn't be changed at all, unless if you run a private Horde Cluster and you know what you are doing.
+backupClusterURL: "https://stablehorde.net"
 
 # An amazing, cool, awesome Horde Worker name you can specify yourself.
 workerName: ""
 
+# The worker type you must specify. If left blank, it'll auto default to text.
+# Acceptable options are, text or interrogation
+workerType: ""
+
 # The KoboldAI Client URL where it can reach your worker.
+# If you are running an interrogation client, you must use https://github.com/LogicismDev/HordeInterrogateAgent
 # The URL can be a cloudflared, nginx, or localtunnel URL as well that is hosted on other machines (e.g. Colab).
 kaiURL: ""
 
 # The API Key that is used to uniquely identify yourself on the Horde Cluster.
-# If you don't have an API Key, visit https://stablehorde.net/register to grab one.
+# If you don't have an API Key, visit https://aihorde.net/register to grab one.
 apiKey: ""
 
 # The usernames to prioritize generations first.
@@ -22,4 +31,11 @@ apiKey: ""
 # If you have no usernames to prioritize, you must put the following format below.
 # priorityUsernames: []
 # Please note that your username is already prioritized.
-priorityUsernames: []
+priorityUsernames: []
+
+# The interrogation forms to enable for the worker
+interrogationForms:
+  - caption
+  - nsfw
+  - interrogation
+  - strip_background

src/main/java/me/Logicism/JavaHordeBridge/HordeBridge.java

@@ -3,6 +3,7 @@
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.fasterxml.jackson.dataformat.yaml.YAMLFactory;
 import me.Logicism.JavaHordeBridge.console.HordeLogger;
+import me.Logicism.JavaHordeBridge.core.InterrogationGenerator;
 import me.Logicism.JavaHordeBridge.runnables.TextHordeRunnable;
 import me.Logicism.JavaHordeBridge.core.KAIGenerator;
 import org.apache.commons.cli.*;
@@ -16,20 +17,23 @@
 
 public class HordeBridge {
 
-    public static int BRIDGE_VERSION = 2;
+    public static int BRIDGE_VERSION = 3;
     public static String BRIDGE_AGENT = "Java Horde Bridge:" + BRIDGE_VERSION + ":https://github.com/LogicismDev/Java-Horde-Bridge";
 
     public static HordeBridge INSTANCE;
     public static int INTERVAL = 1;
 
     private KAIGenerator kaiGenerator;
+    private InterrogationGenerator interroGenerator;
     private HordeLogger logger;
     private HordeConfig config;
 
+    private static String workerType = "text";
     private static String kaiURL = "http://127.0.0.1:5000";
     private static String kaiName = "Java Horde Bridge Worker #" + new Random().nextInt();
     private static String kaiAPIKey = "0000000000";
-    private static String clusterURL = "https://stablehorde.net";
+    private static String clusterURL = "https://aihorde.net";
+    private static String backupClusterURL = "https://stablehorde.net";
     private static String[] priorityUsernames = new String[0];
 
     private static ExecutorService service = Executors.newSingleThreadExecutor();
@@ -41,18 +45,22 @@ public static void main(String[] args) {
 
         Option interval = Option.builder("i").longOpt("interval").argName("interval").hasArg().desc("The interval to check if there are new generations").numberOfArgs(1).build();
         Option config = Option.builder("c").longOpt("config").argName("config").hasArg().desc("The configuration file to grab the Horde Worker information from").numberOfArgs(1).build();
+        Option worker_type = Option.builder("t").longOpt("worker_type").argName("worker_type").hasArg().desc("The Worker type to specify (text/interrogation)").numberOfArgs(1).build();
         Option kai_url = Option.builder("k").longOpt("kai_url").argName("kai_url").hasArg().desc("The KoboldAI URL to grab generations from").numberOfArgs(1).build();
         Option kai_name = Option.builder("n").longOpt("kai_name").argName("kai_name").hasArg().desc("Set the Horde Worker name").numberOfArgs(1).build();
         Option kai_apikey = Option.builder("a").longOpt("kai_apikey").argName("kai_apikey").hasArg().desc("Set the Horde Worker API Key").numberOfArgs(1).build();
         Option cluster_url = Option.builder("u").longOpt("cluster_url").argName("cluster_url").hasArg().desc("Set the cluster url to grab prompts and to send generations to").numberOfArgs(1).build();
+        Option backup_cluster_url = Option.builder("b").longOpt("backup_cluster_url").argName("backup_cluster_url").hasArg().desc("Set the backup cluster url if the main cluster is down").numberOfArgs(1).build();
         Option priority_usernames = Option.builder("p").longOpt("priority_usernames").argName("priority_usernames").hasArg().valueSeparator(',').desc("The usernames to prioritize generations").build();
 
         options.addOption(interval);
         options.addOption(config);
+        options.addOption(worker_type);
         options.addOption(kai_url);
         options.addOption(kai_name);
         options.addOption(kai_apikey);
         options.addOption(cluster_url);
+        options.addOption(backup_cluster_url);
         options.addOption(priority_usernames);
 
         HelpFormatter helper = new HelpFormatter();
@@ -64,7 +72,7 @@ public static void main(String[] args) {
 
             INSTANCE.getLogger().info("Starting Java Horde Worker " + kaiName);
 
-            service.execute(new TextHordeRunnable(INSTANCE, kaiURL, kaiName, kaiAPIKey, clusterURL, priorityUsernames));
+            service.execute(new TextHordeRunnable(INSTANCE, kaiURL, kaiName, kaiAPIKey, clusterURL, backupClusterURL, priorityUsernames));
         } catch (ParseException e) {
             System.out.println(e.getMessage());
             helper.printHelp("Usage", options);
@@ -84,6 +92,13 @@ public HordeBridge(CommandLine cmd) {
         }
 
         if (!configFile.exists()) {
+            if (cmd.hasOption("worker_type")) {
+                if (!cmd.getOptionValue("worker_type").equals("interrogation") || !cmd.getOptionValue("worker_type").equals("text")) {
+                    workerType = "text";
+                } else {
+                    workerType = cmd.getOptionValue("worker_type");
+                }
+            }
             if (cmd.hasOption("kai_url")) {
                 kaiURL = cmd.getOptionValue("kai_url");
             }
@@ -94,23 +109,33 @@ public HordeBridge(CommandLine cmd) {
                 kaiAPIKey = cmd.getOptionValue("kai_apikey");
             }
             if (cmd.hasOption("kai_url")) {
+                kaiURL = cmd.getOptionValue("kai_url");
+            }
+            if (cmd.hasOption("cluster_url")) {
                 clusterURL = cmd.getOptionValue("cluster_url");
             }
+            if (cmd.hasOption("cluster_url")) {
+                backupClusterURL = cmd.getOptionValue("cluster_url");
+            }
             if (cmd.hasOption("priority_usernames")) {
-                cmd.getOptionValues("priority_usernames");
+                priorityUsernames = cmd.getOptionValues("priority_usernames");
             }
         } else {
             ObjectMapper mapper = new ObjectMapper(new YAMLFactory());
 
             try {
                 config = mapper.readValue(configFile, HordeConfig.class);
 
+                if (config.getWorkerType().equals("interrogation") || config.getWorkerType().equals("text")) {
+                    workerType = config.getWorkerType();
+                }
                 kaiURL = config.getKaiURL();
                 if (!kaiName.isEmpty()) {
                     kaiName = config.getWorkerName();
                 }
                 kaiAPIKey = config.getApiKey();
                 clusterURL = config.getClusterURL();
+                backupClusterURL = config.getBackupClusterURL();
                 priorityUsernames = config.getPriorityUsernames().toArray(new String[0]);
             } catch (IOException e) {
                 logger.error("Couldn't load configuration! " + e.getMessage() + " Stopping Java Horde Bridge");
@@ -129,7 +154,11 @@ public KAIGenerator getGenerator() {
         return kaiGenerator;
     }
 
-    public ExecutorService getService() {
-        return service;
+    public InterrogationGenerator getInterroGenerator() {
+        return interroGenerator;
+    }
+
+    public HordeConfig getConfig() {
+        return config;
     }
 }

src/main/java/me/Logicism/JavaHordeBridge/HordeConfig.java

@@ -6,18 +6,28 @@ public class HordeConfig {
 
     private String clusterURL;
 
+    private String backupClusterURL;
+
     private String workerName;
 
+    private String workerType;
+
     private String kaiURL;
 
     private String apiKey;
 
     private List<String> priorityUsernames;
 
+    private List<String> interrogationForms;
+
     public String getWorkerName() {
         return workerName;
     }
 
+    public String getWorkerType() {
+        return workerType;
+    }
+
     public String getKaiURL() {
         return kaiURL;
     }
@@ -26,11 +36,19 @@ public String getClusterURL() {
         return clusterURL;
     }
 
+    public String getBackupClusterURL() {
+        return backupClusterURL;
+    }
+
     public String getApiKey() {
         return apiKey;
     }
 
     public List<String> getPriorityUsernames() {
         return priorityUsernames;
     }
+
+    public List<String> getInterrogationForms() {
+        return interrogationForms;
+    }
 }