Mintplex-Labs · timothycarambat · Aug 22, 2024 · Aug 22, 2024 · Aug 22, 2024
diff --git a/server/utils/agents/aibitat/plugins/http-socket.js b/server/utils/agents/aibitat/plugins/http-socket.js
@@ -0,0 +1,87 @@
+const chalk = require("chalk");
+const { RetryError } = require("../error");
+const { Telemetry } = require("../../../../models/telemetry");
+
+/**
+ * HTTP Interface plugin for Aibitat to emulate a websocket interface in the agent
+ * framework so we dont have to modify the interface for passing messages and responses
+ * in REST or WSS.
+ */
+const httpSocket = {
+  name: "httpSocket",
+  startupConfig: {
+    params: {
+      handler: {
+        required: true,
+      },
+      muteUserReply: {
+        required: false,
+        default: true,
+      },
+      introspection: {
+        required: false,
+        default: true,
+      },
+    },
+  },
+  plugin: function ({
+    handler,
+    muteUserReply = true, // Do not post messages to "USER" back to frontend.
+    introspection = false, // when enabled will attach socket to Aibitat object with .introspect method which reports status updates to frontend.
+  }) {
+    return {
+      name: this.name,
+      setup(aibitat) {
+        aibitat.onError(async (error) => {
+          if (!!error?.message) {
+            console.error(chalk.red(`   error: ${error.message}`), error);
+            aibitat.introspect(
+              `Error encountered while running: ${error.message}`
+            );
+          }
+
+          if (error instanceof RetryError) {
+            console.error(chalk.red(`   retrying in 60 seconds...`));
+            setTimeout(() => {
+              aibitat.retry();
+            }, 60_000);
+            return;
+          }
+        });
+
+        aibitat.introspect = (messageText) => {
+          if (!introspection) return; // Dump thoughts when not wanted.
+          handler.send(
+            JSON.stringify({ type: "statusResponse", content: messageText })
+          );
+        };
+
+        // expose function for sockets across aibitat
+        // type param must be set or else msg will not be shown or handled in UI.
+        aibitat.socket = {
+          send: (type = "__unhandled", content = "") => {
+            handler.send(JSON.stringify({ type, content }));
+          },
+        };
+
+        // We can only receive one message response with HTTP
+        // so we end on first response.
+        aibitat.onMessage((message) => {
+          if (message.from !== "USER")
+            Telemetry.sendTelemetry("agent_chat_sent");
+          if (message.from === "USER" && muteUserReply) return;
+          handler.send(JSON.stringify(message));
+          handler.close();
+        });
+
+        aibitat.onTerminate(() => {
+          handler.close();
+        });
+      },
+    };
+  },
+};
+
+module.exports = {
+  httpSocket,
+};
diff --git a/server/utils/agents/defaults.js b/server/utils/agents/defaults.js
@@ -22,36 +22,48 @@ const WORKSPACE_AGENT = {
       AgentPlugins.webScraping.name, // Collector web-scraping
     ];
 
-    const _setting = (
-      await SystemSettings.get({ label: "default_agent_skills" })
-    )?.value;
-
-    safeJsonParse(_setting, []).forEach((skillName) => {
-      if (!AgentPlugins.hasOwnProperty(skillName)) return;
-
-      // This is a plugin module with many sub-children plugins who
-      // need to be named via `${parent}#${child}` naming convention
-      if (Array.isArray(AgentPlugins[skillName].plugin)) {
-        for (const subPlugin of AgentPlugins[skillName].plugin) {
-          defaultFunctions.push(
-            `${AgentPlugins[skillName].name}#${subPlugin.name}`
-          );
-        }
-        return;
-      }
-
-      // This is normal single-stage plugin
-      defaultFunctions.push(AgentPlugins[skillName].name);
-    });
-
     return {
       role: Provider.systemPrompt(provider),
-      functions: defaultFunctions,
+      functions: [
+        ...defaultFunctions,
+        ...(await agentSkillsFromSystemSettings()),
+      ],
     };
   },
 };
 
+/**
+ * Fetches and preloads the names/identifiers for plugins that will be dynamically
+ * loaded later
+ * @returns {Promise<string[]>}
+ */
+async function agentSkillsFromSystemSettings() {
+  const systemFunctions = [];
+  const _setting = (await SystemSettings.get({ label: "default_agent_skills" }))
+    ?.value;
+
+  safeJsonParse(_setting, []).forEach((skillName) => {
+    if (!AgentPlugins.hasOwnProperty(skillName)) return;
+
+    // This is a plugin module with many sub-children plugins who
+    // need to be named via `${parent}#${child}` naming convention
+    if (Array.isArray(AgentPlugins[skillName].plugin)) {
+      for (const subPlugin of AgentPlugins[skillName].plugin) {
+        systemFunctions.push(
+          `${AgentPlugins[skillName].name}#${subPlugin.name}`
+        );
+      }
+      return;
+    }
+
+    // This is normal single-stage plugin
+    systemFunctions.push(AgentPlugins[skillName].name);
+  });
+  return systemFunctions;
+}
+
 module.exports = {
   USER_AGENT,
   WORKSPACE_AGENT,
+  agentSkillsFromSystemSettings,
 };