azure-sql-ai-assistant
Version:
AI-powered assistant for Azure SQL Database that combines schema inspection, natural language queries, and intelligent result formatting
372 lines (370 loc) • 14.5 kB
JavaScript
;
var __createBinding = (this && this.__createBinding) || (Object.create ? (function(o, m, k, k2) {
if (k2 === undefined) k2 = k;
var desc = Object.getOwnPropertyDescriptor(m, k);
if (!desc || ("get" in desc ? !m.__esModule : desc.writable || desc.configurable)) {
desc = { enumerable: true, get: function() { return m[k]; } };
}
Object.defineProperty(o, k2, desc);
}) : (function(o, m, k, k2) {
if (k2 === undefined) k2 = k;
o[k2] = m[k];
}));
var __setModuleDefault = (this && this.__setModuleDefault) || (Object.create ? (function(o, v) {
Object.defineProperty(o, "default", { enumerable: true, value: v });
}) : function(o, v) {
o["default"] = v;
});
var __importStar = (this && this.__importStar) || function (mod) {
if (mod && mod.__esModule) return mod;
var result = {};
if (mod != null) for (var k in mod) if (k !== "default" && Object.prototype.hasOwnProperty.call(mod, k)) __createBinding(result, mod, k);
__setModuleDefault(result, mod);
return result;
};
Object.defineProperty(exports, "__esModule", { value: true });
exports.SQLAnalyzer = exports.SQLQueryAssistant = exports.SQLSchemaInspector = void 0;
const sql = __importStar(require("mssql"));
// Main classes
/**
* Inspects and analyzes SQL Server database schema.
* Provides detailed information about tables, columns, keys, and indexes.
*/
class SQLSchemaInspector {
/**
* Creates an instance of SQLSchemaInspector.
* @param config - SQL Server connection configuration
*/
constructor(config) {
this.pool = null;
this.retryConfig = {
maxRetries: 3,
delay: 1000,
};
this.config = config;
}
/**
* Establishes connection to the database with retry mechanism.
* @throws Error if connection fails after maximum retry attempts
*/
async connect() {
let attempts = 0;
while (attempts < this.retryConfig.maxRetries) {
try {
if (!this.pool) {
this.pool = await new sql.ConnectionPool(this.config).connect();
}
return;
}
catch (err) {
attempts++;
if (attempts === this.retryConfig.maxRetries) {
throw new Error(`Failed to connect to database after ${attempts} attempts: ${err instanceof Error ? err.message : "Unknown error"}`);
}
await new Promise((resolve) => setTimeout(resolve, this.retryConfig.delay));
}
}
}
/**
* Safely closes the database connection.
*/
async disconnect() {
if (this.pool) {
await this.pool.close();
this.pool = null;
}
}
/**
* Executes a SQL query and returns the results.
* @param query - SQL query to execute
* @returns Promise resolving to an array of query results
* @throws Error if query execution fails
*/
async executeQuery(query) {
if (!this.pool) {
await this.connect();
}
const result = await this.pool.request().query(query);
return result.recordset;
}
/**
* Performs comprehensive analysis of the database schema.
* @returns Promise resolving to a formatted string containing schema details
* @throws Error if schema inspection fails
*/
async inspectSchema() {
try {
await this.connect();
if (!this.pool) {
throw new Error("Failed to establish database connection");
}
// Get all tables
const tables = await this.executeQuery(`
SELECT
t.TABLE_NAME,
CAST(p.value AS NVARCHAR(MAX)) as TABLE_DESCRIPTION
FROM
INFORMATION_SCHEMA.TABLES t
LEFT JOIN sys.extended_properties p ON
p.major_id = OBJECT_ID(t.TABLE_NAME)
AND p.minor_id = 0
AND p.name = 'MS_Description'
WHERE
t.TABLE_TYPE = 'BASE TABLE'
`);
let schema = "Database Schema:\n=================\n\n";
// Process each table
for (const table of tables) {
const tableName = table.TABLE_NAME;
// Get columns information
const columns = await this.executeQuery(`
SELECT
c.COLUMN_NAME,
c.DATA_TYPE,
c.IS_NULLABLE,
c.COLUMN_DEFAULT,
CAST(ep.value AS NVARCHAR(MAX)) as COLUMN_DESCRIPTION
FROM
INFORMATION_SCHEMA.COLUMNS c
LEFT JOIN sys.columns sc ON
sc.object_id = OBJECT_ID(c.TABLE_NAME)
AND sc.name = c.COLUMN_NAME
LEFT JOIN sys.extended_properties ep ON
ep.major_id = sc.object_id
AND ep.minor_id = sc.column_id
AND ep.name = 'MS_Description'
WHERE
c.TABLE_NAME = '${tableName}'
ORDER BY
c.ORDINAL_POSITION
`);
// Add table header with description if exists
schema += `Table: ${tableName}\n`;
if (table.TABLE_DESCRIPTION) {
schema += `Description: ${table.TABLE_DESCRIPTION}\n`;
}
schema += "----------------------------------------\n";
// Add columns information
schema += "Columns:\n";
for (const column of columns) {
schema += `- ${column.COLUMN_NAME} (${column.DATA_TYPE})`;
// Add additional column properties
const properties = [];
if (column.IS_NULLABLE === "NO")
properties.push("NOT NULL");
if (column.COLUMN_DEFAULT !== null)
properties.push(`DEFAULT: ${column.COLUMN_DEFAULT}`);
if (properties.length > 0) {
schema += ` [${properties.join(", ")}]`;
}
// Add column description if exists
if (column.COLUMN_DESCRIPTION) {
schema += `\n Description: ${column.COLUMN_DESCRIPTION}`;
}
schema += "\n";
}
// Get primary keys
const primaryKeys = await this.executeQuery(`
SELECT
COLUMN_NAME
FROM
INFORMATION_SCHEMA.KEY_COLUMN_USAGE
WHERE
OBJECTPROPERTY(OBJECT_ID(CONSTRAINT_NAME), 'IsPrimaryKey') = 1
AND TABLE_NAME = '${tableName}'
`);
if (primaryKeys.length > 0) {
schema += "\nPrimary Keys:\n";
primaryKeys.forEach((pk) => {
schema += `- ${pk.COLUMN_NAME}\n`;
});
}
// Get foreign keys
const foreignKeys = await this.executeQuery(`
SELECT
fk.name as FK_NAME,
OBJECT_NAME(fk.parent_object_id) as TABLE_NAME,
c1.name as COLUMN_NAME,
OBJECT_NAME(fk.referenced_object_id) as REFERENCED_TABLE_NAME,
c2.name as REFERENCED_COLUMN_NAME
FROM
sys.foreign_keys fk
INNER JOIN sys.foreign_key_columns fkc ON fk.object_id = fkc.constraint_object_id
INNER JOIN sys.columns c1 ON fkc.parent_object_id = c1.object_id AND fkc.parent_column_id = c1.column_id
INNER JOIN sys.columns c2 ON fkc.referenced_object_id = c2.object_id AND fkc.referenced_column_id = c2.column_id
WHERE
OBJECT_NAME(fk.parent_object_id) = '${tableName}'
`);
if (foreignKeys.length > 0) {
schema += "\nForeign Keys:\n";
foreignKeys.forEach((fk) => {
schema += `- ${fk.COLUMN_NAME} -> ${fk.REFERENCED_TABLE_NAME}(${fk.REFERENCED_COLUMN_NAME})\n`;
});
}
// Get indexes
const indexes = await this.executeQuery(`
SELECT
i.name as INDEX_NAME,
COL_NAME(ic.object_id, ic.column_id) as COLUMN_NAME,
i.is_unique
FROM
sys.indexes i
INNER JOIN sys.index_columns ic ON i.object_id = ic.object_id AND i.index_id = ic.index_id
WHERE
i.object_id = OBJECT_ID('${tableName}')
AND i.is_primary_key = 0
ORDER BY
i.name, ic.key_ordinal
`);
if (indexes.length > 0) {
schema += "\nIndexes:\n";
const uniqueIndexes = new Map();
indexes.forEach((index) => {
if (!uniqueIndexes.has(index.INDEX_NAME)) {
uniqueIndexes.set(index.INDEX_NAME, {
columns: [],
isUnique: index.is_unique,
});
}
uniqueIndexes
.get(index.INDEX_NAME)
.columns.push(index.COLUMN_NAME);
});
uniqueIndexes.forEach((value, indexName) => {
schema += `- ${indexName}${value.isUnique ? " (UNIQUE)" : ""}: ${value.columns.join(", ")}\n`;
});
}
schema += "\n\n";
}
return schema;
}
catch (error) {
throw new Error(`Failed to inspect database schema: ${error instanceof Error ? error.message : "Unknown error"}`);
}
}
}
exports.SQLSchemaInspector = SQLSchemaInspector;
/**
* Handles AI-powered SQL query generation and result formatting.
*/
class SQLQueryAssistant {
/**
* Creates an instance of SQLQueryAssistant.
* @param config - Configuration for AI client including model and language settings
*/
constructor(config) {
this.aiConfig = {
...config,
language: config.language || "english",
};
}
/**
* Generates SQL query based on natural language prompt and schema.
* @param prompt - Natural language query description
* @param schema - Database schema information
* @returns Promise resolving to generated SQL query
*/
async generateSQLQuery(prompt, schema) {
const response = await this.aiConfig.client.chat.completions.create({
model: this.aiConfig.model,
messages: [
{
role: "system",
content: `You are a SQL expert. Given a database schema, generate a SQL query to answer the user's question.
Database Schema:
${schema}
Rules:
1. Return ONLY the SQL query, nothing else
2. Use proper SQL syntax
3. Make sure the query is safe and efficient`,
},
{
role: "user",
content: prompt,
},
],
temperature: 0,
});
return this.cleanQuery(response.choices[0].message.content || "");
}
/**
* Formats SQL query results into natural language response.
* @param query - Executed SQL query
* @param results - Query execution results
* @returns Promise resolving to formatted natural language response
*/
async formatResponse(query, results) {
const response = await this.aiConfig.client.chat.completions.create({
model: this.aiConfig.model,
messages: [
{
role: "system",
content: `Format the SQL query results into a natural language response in ${this.aiConfig.language}.`,
},
{
role: "user",
content: `SQL Query: ${query}\nResults: ${JSON.stringify(results)}\n\nPlease provide a natural language response:`,
},
],
temperature: 0,
});
return response.choices[0].message.content || "";
}
/**
* Cleans and formats the generated SQL query.
* @param query - Raw query string to clean
* @returns Cleaned query string
* @private
*/
cleanQuery(query) {
return query
.replace(/```sql\n?/g, "")
.replace(/```/g, "")
.replace(/^\s+|\s+$/g, "");
}
}
exports.SQLQueryAssistant = SQLQueryAssistant;
/**
* Main class that combines schema inspection and AI query capabilities.
* Provides end-to-end functionality for natural language database querying.
*/
class SQLAnalyzer {
/**
* Creates an instance of SQLAnalyzer.
* @param dbConfig - SQL Server connection configuration
* @param aiConfig - AI client configuration
*/
constructor(dbConfig, aiConfig) {
this.inspector = new SQLSchemaInspector(dbConfig);
this.assistant = new SQLQueryAssistant(aiConfig);
}
/**
* Processes natural language query through the entire pipeline.
* @param prompt - Natural language query
* @returns Promise resolving to formatted query results
* @throws Error if any step in the process fails
*/
async analyzeAndQuery(prompt) {
try {
// Отримуємо схему бази даних
const schema = await this.inspector.inspectSchema();
// Генеруємо SQL запит на основі промпту
const query = await this.assistant.generateSQLQuery(prompt, schema);
// Виконуємо згенерований запит
const results = await this.inspector.executeQuery(query);
console.log("SQL Results: ", results);
// Форматуємо результати у природню мову
return await this.assistant.formatResponse(query, results);
}
catch (error) {
throw new Error(`Failed to analyze and query: ${error instanceof Error ? error.message : "Unknown error"}`);
}
}
/**
* Closes the database connection.
*/
async disconnect() {
await this.inspector.disconnect();
}
}
exports.SQLAnalyzer = SQLAnalyzer;