totem-llm 0.1.0

package/server/utils/agents/aibitat/plugins/sql-agent/SQLConnectors/index.js

@@ -0,0 +1,80 @@
+const { SystemSettings } = require("../../../../../../models/systemSettings");
+const { safeJsonParse } = require("../../../../../http");
+
+/**
+ * @typedef {('postgresql'|'mysql'|'sql-server')} SQLEngine
+ */
+
+/**
+ * @typedef {Object} QueryResult
+ * @property {[number]} rows - The query result rows
+ * @property {number} count - Number of rows the query returned/changed
+ * @property {string|null} error - Error string if there was an issue
+ */
+
+/**
+ * A valid database SQL connection object
+ * @typedef {Object} SQLConnection
+ * @property {string} database_id - Unique identifier of the database connection
+ * @property {SQLEngine} engine - Engine used by connection
+ * @property {string} connectionString - RFC connection string for db
+ */
+
+/**
+ * @param {SQLEngine} identifier
+ * @param {object} connectionConfig
+ * @returns Database Connection Engine Class for SQLAgent or throws error
+ */
+function getDBClient(identifier = "", connectionConfig = {}) {
+  switch (identifier) {
+    case "mysql":
+      const { MySQLConnector } = require("./MySQL");
+      return new MySQLConnector(connectionConfig);
+    case "postgresql":
+      const { PostgresSQLConnector } = require("./Postgresql");
+      return new PostgresSQLConnector(connectionConfig);
+    case "sql-server":
+      const { MSSQLConnector } = require("./MSSQL");
+      return new MSSQLConnector(connectionConfig);
+    default:
+      throw new Error(
+        `There is no supported database connector for ${identifier}`
+      );
+  }
+}
+
+/**
+ * Lists all of the known database connection that can be used by the agent.
+ * @returns {Promise<[SQLConnection]>}
+ */
+async function listSQLConnections() {
+  return safeJsonParse(
+    (await SystemSettings.get({ label: "agent_sql_connections" }))?.value,
+    []
+  );
+}
+
+/**
+ * Validates a SQL connection by attempting to connect and run a simple query
+ * @param {SQLEngine} identifier - The SQL engine type
+ * @param {object} connectionConfig - The connection configuration
+ * @returns {Promise<{success: boolean, error: string|null}>}
+ */
+async function validateConnection(identifier = "", connectionConfig = {}) {
+  try {
+    const client = getDBClient(identifier, connectionConfig);
+    return await client.validateConnection();
+  } catch {
+    console.log(`Failed to connect to ${identifier} database.`);
+    return {
+      success: false,
+      error: `Unable to connect to ${identifier}. Please verify your connection details.`,
+    };
+  }
+}
+
+module.exports = {
+  getDBClient,
+  listSQLConnections,
+  validateConnection,
+};

package/server/utils/agents/aibitat/plugins/sql-agent/SQLConnectors/utils.js

@@ -0,0 +1,182 @@
+// Credit: https://github.com/sindilevich/connection-string-parser
+
+/**
+ * @typedef {Object} ConnectionStringParserOptions
+ * @property {'mssql' | 'mysql' | 'postgresql' | 'db'} [scheme] - The scheme of the connection string
+ */
+
+/**
+ * @typedef {Object} ConnectionStringObject
+ * @property {string} scheme - The scheme of the connection string eg: mongodb, mssql, mysql, postgresql, etc.
+ * @property {string} username - The username of the connection string
+ * @property {string} password - The password of the connection string
+ * @property {{host: string, port: number}[]} hosts - The hosts of the connection string
+ * @property {string} endpoint - The endpoint (database name) of the connection string
+ * @property {Object} options - The options of the connection string
+ */
+class ConnectionStringParser {
+  static DEFAULT_SCHEME = "db";
+
+  /**
+   * @param {ConnectionStringParserOptions} options
+   */
+  constructor(options = {}) {
+    this.scheme =
+      (options && options.scheme) || ConnectionStringParser.DEFAULT_SCHEME;
+  }
+
+  /**
+   * Takes a connection string object and returns a URI string of the form:
+   *
+   * scheme://[username[:password]@]host1[:port1][,host2[:port2],...[,hostN[:portN]]][/[endpoint]][?options]
+   * @param {Object} connectionStringObject The object that describes connection string parameters
+   */
+  format(connectionStringObject) {
+    if (!connectionStringObject) {
+      return this.scheme + "://localhost";
+    }
+    if (
+      this.scheme &&
+      connectionStringObject.scheme &&
+      this.scheme !== connectionStringObject.scheme
+    ) {
+      throw new Error(`Scheme not supported: ${connectionStringObject.scheme}`);
+    }
+
+    let uri =
+      (this.scheme ||
+        connectionStringObject.scheme ||
+        ConnectionStringParser.DEFAULT_SCHEME) + "://";
+
+    if (connectionStringObject.username) {
+      uri += encodeURIComponent(connectionStringObject.username);
+      // Allow empty passwords
+      if (connectionStringObject.password) {
+        uri += ":" + encodeURIComponent(connectionStringObject.password);
+      }
+      uri += "@";
+    }
+    uri += this._formatAddress(connectionStringObject);
+    // Only put a slash when there is an endpoint
+    if (connectionStringObject.endpoint) {
+      uri += "/" + encodeURIComponent(connectionStringObject.endpoint);
+    }
+    if (
+      connectionStringObject.options &&
+      Object.keys(connectionStringObject.options).length > 0
+    ) {
+      uri +=
+        "?" +
+        Object.keys(connectionStringObject.options)
+          .map(
+            (option) =>
+              encodeURIComponent(option) +
+              "=" +
+              encodeURIComponent(connectionStringObject.options[option])
+          )
+          .join("&");
+    }
+    return uri;
+  }
+
+  /**
+   * Where scheme and hosts will always be present. Other fields will only be present in the result if they were
+   * present in the input.
+   * @param {string} uri The connection string URI
+   * @returns {ConnectionStringObject} The connection string object
+   */
+  parse(uri) {
+    const connectionStringParser = new RegExp(
+      "^\\s*" + // Optional whitespace padding at the beginning of the line
+        "([^:]+)://" + // Scheme (Group 1)
+        "(?:([^:@,/?=&]+)(?::([^:@,/?=&]+))?@)?" + // User (Group 2) and Password (Group 3)
+        "([^@/?=&]+)" + // Host address(es) (Group 4)
+        "(?:/([^:@,/?=&]+)?)?" + // Endpoint (Group 5)
+        "(?:\\?([^:@,/?]+)?)?" + // Options (Group 6)
+        "\\s*$", // Optional whitespace padding at the end of the line
+      "gi"
+    );
+    const connectionStringObject = {};
+
+    if (!uri.includes("://")) {
+      throw new Error(`No scheme found in URI ${uri}`);
+    }
+
+    const tokens = connectionStringParser.exec(uri);
+
+    if (Array.isArray(tokens)) {
+      connectionStringObject.scheme = tokens[1];
+      if (this.scheme && this.scheme !== connectionStringObject.scheme) {
+        throw new Error(`URI must start with '${this.scheme}://'`);
+      }
+      connectionStringObject.username = tokens[2]
+        ? decodeURIComponent(tokens[2])
+        : tokens[2];
+      connectionStringObject.password = tokens[3]
+        ? decodeURIComponent(tokens[3])
+        : tokens[3];
+      connectionStringObject.hosts = this._parseAddress(tokens[4]);
+      connectionStringObject.endpoint = tokens[5]
+        ? decodeURIComponent(tokens[5])
+        : tokens[5];
+      connectionStringObject.options = tokens[6]
+        ? this._parseOptions(tokens[6])
+        : tokens[6];
+    }
+    return connectionStringObject;
+  }
+
+  /**
+   * Formats the address portion of a connection string
+   * @param {Object} connectionStringObject The object that describes connection string parameters
+   */
+  _formatAddress(connectionStringObject) {
+    return connectionStringObject.hosts
+      .map(
+        (address) =>
+          encodeURIComponent(address.host) +
+          (address.port
+            ? ":" + encodeURIComponent(address.port.toString(10))
+            : "")
+      )
+      .join(",");
+  }
+
+  /**
+   * Parses an address
+   * @param {string} addresses The address(es) to process
+   */
+  _parseAddress(addresses) {
+    return addresses.split(",").map((address) => {
+      const i = address.indexOf(":");
+
+      return i >= 0
+        ? {
+            host: decodeURIComponent(address.substring(0, i)),
+            port: +address.substring(i + 1),
+          }
+        : { host: decodeURIComponent(address) };
+    });
+  }
+
+  /**
+   * Parses options
+   * @param {string} options The options to process
+   */
+  _parseOptions(options) {
+    const result = {};
+
+    options.split("&").forEach((option) => {
+      const i = option.indexOf("=");
+
+      if (i >= 0) {
+        result[decodeURIComponent(option.substring(0, i))] = decodeURIComponent(
+          option.substring(i + 1)
+        );
+      }
+    });
+    return result;
+  }
+}
+
+module.exports = { ConnectionStringParser };

package/server/utils/agents/aibitat/plugins/sql-agent/get-table-schema.js

@@ -0,0 +1,115 @@
+module.exports.SqlAgentGetTableSchema = {
+  name: "sql-get-table-schema",
+  plugin: function () {
+    const {
+      listSQLConnections,
+      getDBClient,
+    } = require("./SQLConnectors/index.js");
+
+    function formatQueryForDisplay(query, params = []) {
+      if (!params.length) return query;
+      let formatted = query;
+      params.forEach((param, index) => {
+        const value = typeof param === "string" ? `'${param}'` : param;
+        formatted = formatted.replace(`$${index + 1}`, value);
+        formatted = formatted.replace(`@p${index}`, value);
+        formatted = formatted.replace("?", value);
+      });
+      return formatted;
+    }
+
+    return {
+      name: "sql-get-table-schema",
+      setup(aibitat) {
+        aibitat.function({
+          super: aibitat,
+          name: this.name,
+          description:
+            "Gets the table schema in SQL for a given `table` and `database_id`",
+          examples: [
+            {
+              prompt: "What does the customers table in access-logs look like?",
+              call: JSON.stringify({
+                database_id: "access-logs",
+                table_name: "customers",
+              }),
+            },
+            {
+              prompt:
+                "Get me the full name of a company in records-main, the table should be call comps",
+              call: JSON.stringify({
+                database_id: "records-main",
+                table_name: "comps",
+              }),
+            },
+          ],
+          parameters: {
+            $schema: "http://json-schema.org/draft-07/schema#",
+            type: "object",
+            properties: {
+              database_id: {
+                type: "string",
+                description:
+                  "The database identifier for which we will connect to to query the table schema. This is a required field.",
+              },
+              table_name: {
+                type: "string",
+                description:
+                  "The database identifier for the table name we want the schema for. This is a required field.",
+              },
+            },
+            additionalProperties: false,
+          },
+          required: ["database_id", "table_name"],
+          handler: async function ({ database_id = "", table_name = "" }) {
+            this.super.handlerProps.log(`Using the sql-get-table-schema tool.`);
+            try {
+              const databaseConfig = (await listSQLConnections()).find(
+                (db) => db.database_id === database_id
+              );
+              if (!databaseConfig) {
+                this.super.handlerProps.log(
+                  `sql-get-table-schema to find config!`,
+                  database_id
+                );
+                return `No database connection for ${database_id} was found!`;
+              }
+
+              const db = getDBClient(databaseConfig.engine, databaseConfig);
+              this.super.introspect(
+                `${this.caller}: Querying the table schema for ${table_name} in the ${databaseConfig.database_id} database.`
+              );
+
+              const sqlQuery = db.getTableSchemaSql(table_name);
+              const isParameterized =
+                typeof sqlQuery === "object" && sqlQuery.query;
+              const queryString = isParameterized ? sqlQuery.query : sqlQuery;
+              const queryParams = isParameterized ? sqlQuery.params : [];
+
+              this.super.introspect(
+                `Running SQL: ${formatQueryForDisplay(queryString, queryParams)}`
+              );
+              const result = await db.runQuery(queryString, queryParams);
+
+              if (result.error) {
+                this.super.handlerProps.log(
+                  `sql-get-table-schema tool reported error`,
+                  result.error
+                );
+                this.super.introspect(`Error: ${result.error}`);
+                return `There was an error running the query: ${result.error}`;
+              }
+
+              return JSON.stringify(result);
+            } catch (e) {
+              this.super.handlerProps.log(
+                `sql-get-table-schema raised an error. ${e.message}`
+              );
+              return e.message;
+            }
+          },
+        });
+      },
+    };
+  },
+};

package/server/utils/agents/aibitat/plugins/sql-agent/index.js

@@ -0,0 +1,21 @@
+const { SqlAgentGetTableSchema } = require("./get-table-schema");
+const { SqlAgentListDatabase } = require("./list-database");
+const { SqlAgentListTables } = require("./list-table");
+const { SqlAgentQuery } = require("./query");
+
+const sqlAgent = {
+  name: "sql-agent",
+  startupConfig: {
+    params: {},
+  },
+  plugin: [
+    SqlAgentListDatabase,
+    SqlAgentListTables,
+    SqlAgentGetTableSchema,
+    SqlAgentQuery,
+  ],
+};
+
+module.exports = {
+  sqlAgent,
+};

package/server/utils/agents/aibitat/plugins/sql-agent/list-database.js

@@ -0,0 +1,49 @@
+module.exports.SqlAgentListDatabase = {
+  name: "sql-list-databases",
+  plugin: function () {
+    const { listSQLConnections } = require("./SQLConnectors");
+    return {
+      name: "sql-list-databases",
+      setup(aibitat) {
+        aibitat.function({
+          super: aibitat,
+          name: this.name,
+          description:
+            "List all available databases via `list_databases` you currently have access to. Returns a unique string identifier `database_id` that can be used for future calls.",
+          examples: [
+            {
+              prompt: "What databases can you access?",
+              call: JSON.stringify({}),
+            },
+            {
+              prompt: "What databases can you tell me about?",
+              call: JSON.stringify({}),
+            },
+            {
+              prompt: "Is there a database named erp-logs you can access?",
+              call: JSON.stringify({}),
+            },
+          ],
+          parameters: {
+            $schema: "http://json-schema.org/draft-07/schema#",
+            type: "object",
+            properties: {},
+            additionalProperties: false,
+          },
+          handler: async function () {
+            this.super.handlerProps.log(`Using the sql-list-databases tool.`);
+            this.super.introspect(
+              `${this.caller}: Checking what are the available databases.`
+            );
+
+            const connections = (await listSQLConnections()).map((conn) => {
+              const { connectionString: _connectionString, ...rest } = conn;
+              return rest;
+            });
+            return JSON.stringify(connections);
+          },
+        });
+      },
+    };
+  },
+};

package/server/utils/agents/aibitat/plugins/sql-agent/list-table.js

@@ -0,0 +1,105 @@
+module.exports.SqlAgentListTables = {
+  name: "sql-list-tables",
+  plugin: function () {
+    const {
+      listSQLConnections,
+      getDBClient,
+    } = require("./SQLConnectors/index.js");
+
+    function formatQueryForDisplay(query, params = []) {
+      if (!params.length) return query;
+      let formatted = query;
+      params.forEach((param, index) => {
+        const value = typeof param === "string" ? `'${param}'` : param;
+        formatted = formatted.replace(`$${index + 1}`, value);
+        formatted = formatted.replace(`@p${index}`, value);
+        formatted = formatted.replace("?", value);
+      });
+      return formatted;
+    }
+
+    return {
+      name: "sql-list-tables",
+      setup(aibitat) {
+        aibitat.function({
+          super: aibitat,
+          name: this.name,
+          description:
+            "List all available tables in a database via its `database_id`.",
+          examples: [
+            {
+              prompt: "What tables are there in the `access-logs` database?",
+              call: JSON.stringify({ database_id: "access-logs" }),
+            },
+            {
+              prompt:
+                "What information can you access in the customer_accts postgres db?",
+              call: JSON.stringify({ database_id: "customer_accts" }),
+            },
+            {
+              prompt: "Can you tell me what is in the primary-logs db?",
+              call: JSON.stringify({ database_id: "primary-logs" }),
+            },
+          ],
+          parameters: {
+            $schema: "http://json-schema.org/draft-07/schema#",
+            type: "object",
+            properties: {
+              database_id: {
+                type: "string",
+                description:
+                  "The database identifier for which we will list all tables for. This is a required parameter",
+              },
+            },
+            additionalProperties: false,
+          },
+          required: ["database_id"],
+          handler: async function ({ database_id = "" }) {
+            try {
+              this.super.handlerProps.log(`Using the sql-list-tables tool.`);
+              const databaseConfig = (await listSQLConnections()).find(
+                (db) => db.database_id === database_id
+              );
+              if (!databaseConfig) {
+                this.super.handlerProps.log(
+                  `sql-list-tables failed to find config!`,
+                  database_id
+                );
+                return `No database connection for ${database_id} was found!`;
+              }
+
+              const db = getDBClient(databaseConfig.engine, databaseConfig);
+              this.super.introspect(
+                `${this.caller}: Checking what are the available tables in the ${databaseConfig.database_id} database.`
+              );
+
+              const sqlQuery = db.getTablesSql();
+              const isParameterized =
+                typeof sqlQuery === "object" && sqlQuery.query;
+              const queryString = isParameterized ? sqlQuery.query : sqlQuery;
+              const queryParams = isParameterized ? sqlQuery.params : [];
+
+              this.super.introspect(
+                `Running SQL: ${formatQueryForDisplay(queryString, queryParams)}`
+              );
+              const result = await db.runQuery(queryString, queryParams);
+              if (result.error) {
+                this.super.handlerProps.log(
+                  `sql-list-tables tool reported error`,
+                  result.error
+                );
+                this.super.introspect(`Error: ${result.error}`);
+                return `There was an error running the query: ${result.error}`;
+              }
+
+              return JSON.stringify(result);
+            } catch (e) {
+              console.error(e);
+              return e.message;
+            }
+          },
+        });
+      },
+    };
+  },
+};

package/server/utils/agents/aibitat/plugins/sql-agent/query.js

@@ -0,0 +1,101 @@
+module.exports.SqlAgentQuery = {
+  name: "sql-query",
+  plugin: function () {
+    const {
+      getDBClient,
+      listSQLConnections,
+    } = require("./SQLConnectors/index.js");
+
+    return {
+      name: "sql-query",
+      setup(aibitat) {
+        aibitat.function({
+          super: aibitat,
+          name: this.name,
+          description:
+            "Run a read-only SQL query on a `database_id` which will return up rows of data related to the query. The query must only be SELECT statements which do not modify the table data. There should be a reasonable LIMIT on the return quantity to prevent long-running or queries which crash the db.",
+          examples: [
+            {
+              prompt: "How many customers are in dvd-rentals?",
+              call: JSON.stringify({
+                database_id: "dvd-rentals",
+                sql_query: "SELECT * FROM customers",
+              }),
+            },
+            {
+              prompt: "Can you tell me the total volume of sales last month?",
+              call: JSON.stringify({
+                database_id: "sales-db",
+                sql_query:
+                  "SELECT SUM(sale_amount) AS total_sales FROM sales WHERE sale_date >= DATEADD(month, -1, DATEFROMPARTS(YEAR(GETDATE()), MONTH(GETDATE()), 1)) AND sale_date < DATEFROMPARTS(YEAR(GETDATE()), MONTH(GETDATE()), 1)",
+              }),
+            },
+            {
+              prompt:
+                "Do we have anyone in the staff table for our production db named 'sam'? ",
+              call: JSON.stringify({
+                database_id: "production",
+                sql_query:
+                  "SElECT * FROM staff WHERE first_name='sam%' OR last_name='sam%'",
+              }),
+            },
+          ],
+          parameters: {
+            $schema: "http://json-schema.org/draft-07/schema#",
+            type: "object",
+            properties: {
+              database_id: {
+                type: "string",
+                description:
+                  "The database identifier for which we will connect to to query the table schema. This is required to run the SQL query.",
+              },
+              sql_query: {
+                type: "string",
+                description:
+                  "The raw SQL query to run. Should be a query which does not modify the table and will return results.",
+              },
+            },
+            additionalProperties: false,
+          },
+          required: ["database_id", "table_name"],
+          handler: async function ({ database_id = "", sql_query = "" }) {
+            this.super.handlerProps.log(`Using the sql-query tool.`);
+            try {
+              const databaseConfig = (await listSQLConnections()).find(
+                (db) => db.database_id === database_id
+              );
+              if (!databaseConfig) {
+                this.super.handlerProps.log(
+                  `sql-query failed to find config!`,
+                  database_id
+                );
+                return `No database connection for ${database_id} was found!`;
+              }
+
+              this.super.introspect(
+                `${this.caller}: Im going to run a query on the ${database_id} to get an answer.`
+              );
+              const db = getDBClient(databaseConfig.engine, databaseConfig);
+
+              this.super.introspect(`Running SQL: ${sql_query}`);
+              const result = await db.runQuery(sql_query);
+              if (result.error) {
+                this.super.handlerProps.log(
+                  `sql-query tool reported error`,
+                  result.error
+                );
+                this.super.introspect(`Error: ${result.error}`);
+                return `There was an error running the query: ${result.error}`;
+              }
+
+              return JSON.stringify(result);
+            } catch (e) {
+              console.error(e);
+              return e.message;
+            }
+          },
+        });
+      },
+    };
+  },
+};