Added postgres, cockroachdb, mssql, mysql, mariadb, mongodb and oracle to SqlDatabaseChain_Chains

This commit is contained in:
Lorenzo
2023-08-14 00:31:07 +02:00
parent fb115a5f69
commit ff7ee41758
@@ -5,6 +5,9 @@ import { DataSource } from 'typeorm'
import { SqlDatabase } from 'langchain/sql_db' import { SqlDatabase } from 'langchain/sql_db'
import { BaseLanguageModel } from 'langchain/base_language' import { BaseLanguageModel } from 'langchain/base_language'
import { ConsoleCallbackHandler, CustomChainHandler } from '../../../src/handler' import { ConsoleCallbackHandler, CustomChainHandler } from '../../../src/handler'
import { DataSourceOptions } from 'typeorm/data-source'
type DatabaseType = 'sqlite' | 'postgres' | 'cockroachdb' | 'mssql' | 'mysql' | 'mariadb' | 'mongodb' | 'oracle'
class SqlDatabaseChain_Chains implements INode { class SqlDatabaseChain_Chains implements INode {
label: string label: string
@@ -38,36 +41,64 @@ class SqlDatabaseChain_Chains implements INode {
type: 'options', type: 'options',
options: [ options: [
{ {
label: 'SQlite', label: 'SQLite',
name: 'sqlite' name: 'sqlite'
},
{
label: 'PostgreSQL',
name: 'postgres'
},
{
label: 'CockroachDB',
name: 'cockroachdb'
},
{
label: 'MSSQL',
name: 'mssql'
},
{
label: 'MySQL',
name: 'mysql'
},
{
label: 'MariaDB',
name: 'mariadb'
},
{
label: 'MongoDB',
name: 'mongodb'
},
{
label: 'Oracle',
name: 'oracle'
} }
], ],
default: 'sqlite' default: 'sqlite'
}, },
{ {
label: 'Database File Path', label: 'Connection string or file path (sqlite only)',
name: 'dbFilePath', name: 'url',
type: 'string', type: 'string',
placeholder: 'C:/Users/chinook.db' placeholder: '1270.0.0.1:5432/chinook'
} }
] ]
} }
async init(nodeData: INodeData): Promise<any> { async init(nodeData: INodeData): Promise<any> {
const databaseType = nodeData.inputs?.database as 'sqlite' const databaseType = nodeData.inputs?.database as DatabaseType
const model = nodeData.inputs?.model as BaseLanguageModel const model = nodeData.inputs?.model as BaseLanguageModel
const dbFilePath = nodeData.inputs?.dbFilePath const url = nodeData.inputs?.url
const chain = await getSQLDBChain(databaseType, dbFilePath, model) const chain = await getSQLDBChain(databaseType, url, model)
return chain return chain
} }
async run(nodeData: INodeData, input: string, options: ICommonObject): Promise<string> { async run(nodeData: INodeData, input: string, options: ICommonObject): Promise<string> {
const databaseType = nodeData.inputs?.database as 'sqlite' const databaseType = nodeData.inputs?.database as DatabaseType
const model = nodeData.inputs?.model as BaseLanguageModel const model = nodeData.inputs?.model as BaseLanguageModel
const dbFilePath = nodeData.inputs?.dbFilePath const url = nodeData.inputs?.url
const chain = await getSQLDBChain(databaseType, dbFilePath, model) const chain = await getSQLDBChain(databaseType, url, model)
const loggerHandler = new ConsoleCallbackHandler(options.logger) const loggerHandler = new ConsoleCallbackHandler(options.logger)
if (options.socketIO && options.socketIOClientId) { if (options.socketIO && options.socketIOClientId) {
@@ -81,11 +112,18 @@ class SqlDatabaseChain_Chains implements INode {
} }
} }
const getSQLDBChain = async (databaseType: 'sqlite', dbFilePath: string, llm: BaseLanguageModel) => { const getSQLDBChain = async (databaseType: DatabaseType, url: string, llm: BaseLanguageModel) => {
const datasource = new DataSource({ const datasource = new DataSource(
type: databaseType, databaseType === 'sqlite'
database: dbFilePath ? {
}) type: databaseType,
database: url
}
: ({
type: databaseType,
url: url
} as DataSourceOptions)
)
const db = await SqlDatabase.fromDataSourceParams({ const db = await SqlDatabase.fromDataSourceParams({
appDataSource: datasource appDataSource: datasource