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

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;
            }
          },
        });
      },
    };
  },
};