Provide support for multiple schemas for SQL based tools

[Strernd]

Checked

• I searched existing ideas and did not find a similar one
• I added a very descriptive title
• I've clearly described the feature request and motivation for it

Feature request

We have multiple schemas in the database that we use and join. It is currently not possible to pass multiple schemas to the initialization of the database. This seems to be a restriction of typeorm and langchain.

Motivation

The SQL based tools can not be used when the context needs to be aware of tables in different schemas.

Proposal (If applicable)

A not perfect solution would be to load all schemas into the search_path and not use the public schema as a default. A more robust solution would require that typeorm is capable of initializing a client with multiple schemas.
Maybe it's also possible to provide a custom database adapter with a run and introspect method.

[Strernd]

Leaving my hacky workaround here if anyone encounters the same problem:

Step 1. New file with adapted methods from the original langchainjs repo, getting rid of the schema:
Note: I stripped everything that was not postgres

import { SqlTable } from "langchain/dist/util/sql_utils";
import { SqlDatabase } from "langchain/sql_db";
import { DataSource } from "typeorm";

export async function getTableInfo(
  db: SqlDatabase,
  targetTables: Array<string>
): Promise<string> {
  const allTables = await getTableAndColumnsName(db.appDataSource);
  const selectedTables = allTables.filter((currentTable) =>
    targetTables.includes(currentTable.tableName)
  );
  return generateTableInfoFromTables(
    selectedTables,
    db.appDataSource,
    db.sampleRowsInTableInfo,
    db.customDescription
  );
}

const formatToSqlTable = (rawResultsTableAndColumn: any): Array<SqlTable> => {
  const sqlTable: Array<SqlTable> = [];
  for (const oneResult of rawResultsTableAndColumn) {
    const sqlColumn = {
      columnName: oneResult.column_name,
      dataType: oneResult.data_type,
      isNullable: oneResult.is_nullable === "YES",
    };
    const currentTable = sqlTable.find(
      (oneTable) => oneTable.tableName === oneResult.table_name
    );
    if (currentTable) {
      currentTable.columns.push(sqlColumn);
    } else {
      const newTable = {
        tableName: oneResult.table_name,
        columns: [sqlColumn],
      };
      sqlTable.push(newTable);
    }
  }

  return sqlTable;
};

const getTableAndColumnsName = async (
  appDataSource: DataSource
): Promise<Array<SqlTable>> => {
  const sql = `SELECT 
              t.table_name, 
              c.* 
            FROM 
              information_schema.tables t 
                JOIN information_schema.columns c 
                  ON t.table_name = c.table_name 
            ORDER BY 
              t.table_name,
              c.ordinal_position;`;
  const rep = await appDataSource.query(sql);

  return formatToSqlTable(rep);
};

const generateTableInfoFromTables = async (
  tables: Array<SqlTable> | undefined,
  appDataSource: DataSource,
  nbSampleRow: number,
  customDescription?: Record<string, string>
): Promise<string> => {
  if (!tables) {
    return "";
  }

  let globalString = "";
  for (const currentTable of tables) {
    const tableCustomDescription =
      customDescription &&
      Object.keys(customDescription).includes(currentTable.tableName)
        ? `${customDescription[currentTable.tableName]}\n`
        : "";
    let sqlCreateTableQuery = `CREATE TABLE ${currentTable.tableName} (\n`;
    for (const [key, currentColumn] of currentTable.columns.entries()) {
      if (key > 0) {
        sqlCreateTableQuery += ", ";
      }
      sqlCreateTableQuery += `${currentColumn.columnName} ${
        currentColumn.dataType
      } ${currentColumn.isNullable ? "" : "NOT NULL"}`;
    }
    sqlCreateTableQuery += ") \n";

    const sqlSelectInfoQuery = `SELECT * FROM "${currentTable.tableName}" LIMIT ${nbSampleRow};\n`;
    const columnNamesConcatString = `${currentTable.columns.reduce(
      (completeString, column) => `${completeString} ${column.columnName}`,
      ""
    )}\n`;

    let sample = "";
    try {
      const infoObjectResult = nbSampleRow
        ? await appDataSource.query(sqlSelectInfoQuery)
        : null;
      sample = formatSqlResponseToSimpleTableString(infoObjectResult);
    } catch (error) {
      console.log(error);
    }

    globalString = globalString.concat(
      tableCustomDescription +
        sqlCreateTableQuery +
        sqlSelectInfoQuery +
        columnNamesConcatString +
        sample
    );
  }

  return globalString;
};

const formatSqlResponseToSimpleTableString = (rawResult: unknown): string => {
  if (!rawResult || !Array.isArray(rawResult) || rawResult.length === 0) {
    return "";
  }

  let globalString = "";
  for (const oneRow of rawResult) {
    globalString += `${Object.values(oneRow).reduce(
      (completeString, columnValue) => `${completeString} ${columnValue}`,
      ""
    )}\n`;
  }

  return globalString;
};

Step 2: After initializing the database set the search path and overwrite the getTableInfo with your custom one

  await db.run(`SET search_path TO ${schemas.join(",")};`);
  const tableInfo = await getTableInfo(db, tables); // Pass all tables without schema here

  db.getTableInfo = async (ts?: string[]) => {
    return tableInfo;
  };