rag-cli-tester
Version:
A lightweight CLI tool for testing RAG (Retrieval-Augmented Generation) systems with different embedding combinations
315 lines • 12.2 kB
JavaScript
"use strict";
Object.defineProperty(exports, "__esModule", { value: true });
exports.DatabaseConnection = void 0;
const supabase_js_1 = require("@supabase/supabase-js");
class DatabaseConnection {
constructor(config) {
this.isConnected = false;
this.supabase = (0, supabase_js_1.createClient)(config.url, config.anonKey);
}
async testConnection() {
try {
const { data, error } = await this.supabase
.rpc('test_connection');
// Connection is successful if we get data (array of tables) and no error
this.isConnected = !error && data && Array.isArray(data) && data.length > 0;
return this.isConnected;
}
catch (error) {
console.error('Database connection failed:', error);
this.isConnected = false;
return false;
}
}
async getTables() {
try {
const { data, error } = await this.supabase
.rpc('test_connection');
if (error) {
console.error('Failed to fetch tables:', error);
return [];
}
return data?.map((row) => row.table_name) || [];
}
catch (error) {
console.error('Failed to fetch tables:', error);
return [];
}
}
async getTableInfo(tableName) {
try {
// Get column information using RPC function
const { data: columnsData, error: columnsError } = await this.supabase
.rpc('get_table_columns', { table_name_param: tableName });
if (columnsError)
throw columnsError;
if (!columnsData || columnsData.length === 0) {
return null;
}
// Get row count using RPC function
const { data: rowCount, error: countError } = await this.supabase
.rpc('get_table_row_count', { table_name_param: tableName });
if (countError) {
console.warn(`Could not get row count for ${tableName}:`, countError);
}
return {
name: tableName,
columns: columnsData.map((row) => ({
column_name: row.column_name,
data_type: row.data_type,
is_nullable: row.is_nullable === 'YES'
})),
rowCount: rowCount || 0
};
}
catch (error) {
console.error(`Failed to get table info for ${tableName}:`, error);
return null;
}
}
async getTableData(tableName, columns = ['*'], limit, offset) {
try {
let query = this.supabase
.from(tableName)
.select(columns.join(','));
if (limit)
query = query.limit(limit);
if (offset)
query = query.range(offset, offset + (limit || 1000) - 1);
const { data, error } = await query;
if (error)
throw error;
return data || [];
}
catch (error) {
console.error(`Failed to fetch data from ${tableName}:`, error);
return [];
}
}
isConnectionActive() {
return this.isConnected;
}
async getColumnDataType(tableName, columnName) {
try {
const { data, error } = await this.supabase
.rpc('get_column_data_type', {
table_name_param: tableName,
column_name_param: columnName
});
if (error)
throw error;
return data;
}
catch (error) {
console.error(`Failed to get column data type for ${tableName}.${columnName}:`, error);
return null;
}
}
async updateRowEmbedding(tableName, rowId, embeddingColumn, embedding) {
try {
const { error } = await this.supabase
.from(tableName)
.update({ [embeddingColumn]: embedding })
.eq('id', rowId);
if (error)
throw error;
return true;
}
catch (error) {
console.error(`Failed to update embedding for row ${rowId}:`, error);
return false;
}
}
async updateRowColumn(tableName, rowId, columnName, value) {
try {
const { error } = await this.supabase
.from(tableName)
.update({ [columnName]: value })
.eq('id', rowId);
if (error)
throw error;
return true;
}
catch (error) {
console.error(`Failed to update column ${columnName} for row ${rowId}:`, error);
return false;
}
}
async getRowsWithoutEmbeddings(tableName, embeddingColumn, sourceColumns, limit = 100) {
try {
const selectColumns = ['id', ...sourceColumns];
const { data, error } = await this.supabase
.from(tableName)
.select(selectColumns.join(','))
.is(embeddingColumn, null)
.limit(limit);
if (error)
throw error;
return data || [];
}
catch (error) {
console.error(`Failed to get rows without embeddings:`, error);
return [];
}
}
async getRowsWithEmptyColumn(tableName, targetColumn, sourceColumns, limit = 100) {
try {
const selectColumns = ['id', ...sourceColumns, targetColumn];
// Get rows where target column is null
const { data: nullData, error: nullError } = await this.supabase
.from(tableName)
.select(selectColumns.join(','))
.is(targetColumn, null)
.limit(limit);
if (nullError) {
console.error(`Error getting null rows:`, nullError);
return [];
}
// Get rows where target column is empty string
const { data: emptyData, error: emptyError } = await this.supabase
.from(tableName)
.select(selectColumns.join(','))
.eq(targetColumn, '')
.limit(limit);
if (emptyError) {
console.error(`Error getting empty rows:`, emptyError);
return nullData || [];
}
// Combine and deduplicate results
const allRows = [...(nullData || []), ...(emptyData || [])];
const uniqueRows = allRows.filter((row, index, self) => index === self.findIndex(r => r.id === row.id));
return uniqueRows.slice(0, limit);
}
catch (error) {
console.error(`Failed to get rows with empty column:`, error);
return [];
}
}
async checkColumnExists(tableName, columnName) {
try {
const tableInfo = await this.getTableInfo(tableName);
if (!tableInfo)
return false;
return tableInfo.columns.some(col => col.column_name === columnName);
}
catch (error) {
console.error(`Failed to check if column exists:`, error);
return false;
}
}
async getColumnDataCount(tableName, columnName) {
try {
// First try to get count of non-null values
const { count: nonNullCount, error: nonNullError } = await this.supabase
.from(tableName)
.select('*', { count: 'exact', head: true })
.not(columnName, 'is', null);
if (nonNullError) {
console.error(`Error counting non-null values:`, nonNullError);
return 0;
}
// Then try to get count of non-empty string values
const { count: nonEmptyCount, error: nonEmptyError } = await this.supabase
.from(tableName)
.select('*', { count: 'exact', head: true })
.not(columnName, 'eq', '');
if (nonEmptyError) {
console.error(`Error counting non-empty values:`, nonEmptyError);
return nonNullCount || 0;
}
// Return the higher count (non-null count should be >= non-empty count)
return Math.max(nonNullCount || 0, nonEmptyCount || 0);
}
catch (error) {
console.error(`Failed to get column data count:`, error);
return 0;
}
}
async getEmptyColumnCount(tableName, columnName) {
try {
// Count rows where column is null OR empty string
// We need to use separate queries since Supabase doesn't support complex OR with count
const { count: nullCount, error: nullError } = await this.supabase
.from(tableName)
.select('*', { count: 'exact', head: true })
.is(columnName, null);
if (nullError) {
console.error(`Error counting null values:`, nullError);
return 0;
}
const { count: emptyCount, error: emptyError } = await this.supabase
.from(tableName)
.select('*', { count: 'exact', head: true })
.eq(columnName, '');
if (emptyError) {
console.error(`Error counting empty values:`, emptyError);
return nullCount || 0;
}
return (nullCount || 0) + (emptyCount || 0);
}
catch (error) {
console.error(`Failed to get empty column count:`, error);
return 0;
}
}
async getRowColumnValue(tableName, rowId, columnName) {
try {
const { data, error } = await this.supabase
.from(tableName)
.select(columnName)
.eq('id', rowId)
.single();
if (error)
throw error;
return data?.[columnName] || null;
}
catch (error) {
console.error(`Failed to get column value for row ${rowId}:`, error);
return null;
}
}
async getTableDataSample(tableName, sampleSize, ratio) {
try {
// For large datasets, use efficient sampling
// First get total count
const { count: totalCount, error: countError } = await this.supabase
.from(tableName)
.select('*', { count: 'exact', head: true });
if (countError)
throw countError;
if (!totalCount || totalCount <= sampleSize) {
// If dataset is small, return all data
return await this.getTableData(tableName);
}
// Calculate offset based on ratio
const offset = Math.floor(totalCount * ratio);
// Use random sampling with offset for large datasets
// This is more efficient than loading all data into memory
const { data, error } = await this.supabase
.from(tableName)
.select('*')
.range(offset, offset + sampleSize - 1);
if (error)
throw error;
// If we got fewer rows than requested, try to get more
if (data && data.length < sampleSize) {
const remaining = sampleSize - data.length;
const { data: additionalData, error: additionalError } = await this.supabase
.from(tableName)
.select('*')
.range(0, remaining - 1);
if (!additionalError && additionalData) {
data.push(...additionalData);
}
}
return data || [];
}
catch (error) {
console.error(`Failed to get table data sample:`, error);
// Fallback to regular method
return await this.getTableData(tableName, ['*'], sampleSize);
}
}
}
exports.DatabaseConnection = DatabaseConnection;
//# sourceMappingURL=database.js.map