Implement batch processing for AI-generated table settings and enhance error handling

Author: Artuomka

backend/src/entities/ai/ai.service.ts

@@ -27,18 +27,44 @@ interface AIResponse {
 	tables: AIGeneratedTableSettings[];
 }
 
+const AI_BATCH_SIZE = 10;
+
 @Injectable()
 export class AiService {
 	constructor(protected readonly aiCoreService: AICoreService) {}
 
 	public async generateNewTableSettingsWithAI(
 		tablesInformation: Array<TableInformation>,
 	): Promise<Array<TableSettingsEntity>> {
+		const allSettings: Array<TableSettingsEntity> = [];
+
+		for (let i = 0; i < tablesInformation.length; i += AI_BATCH_SIZE) {
+			const batch = tablesInformation.slice(i, i + AI_BATCH_SIZE);
+			try {
+				const batchSettings = await this.processTablesBatch(batch);
+				allSettings.push(...batchSettings);
+			} catch (error) {
+				console.warn(`Batch processing failed, falling back to individual table processing: ${error.message}`);
+				for (const tableInfo of batch) {
+					try {
+						const singleTableSettings = await this.processTablesBatch([tableInfo]);
+						allSettings.push(...singleTableSettings);
+					} catch (singleError) {
+						console.error(`Error processing AI for table "${tableInfo.table_name}": ${singleError.message}`);
+					}
+				}
+			}
+		}
+
+		return allSettings;
+	}
+
+	private async processTablesBatch(tablesInformation: Array<TableInformation>): Promise<Array<TableSettingsEntity>> {
 		const prompt = this.buildPrompt(tablesInformation);
 		const aiResponse = await this.aiCoreService.completeWithProvider(AIProviderType.BEDROCK, prompt, {
 			temperature: 0.3,
 		});
-		const parsedResponse = this.parseAIResponse(aiResponse);
+		const parsedResponse = this.parseAIResponse(aiResponse, tablesInformation);
 		return this.buildTableSettingsEntities(parsedResponse, tablesInformation);
 	}
 
@@ -131,13 +157,14 @@ Respond ONLY with valid JSON in this exact format (no markdown, no explanations)
 }`;
 	}
 
-	private parseAIResponse(aiResponse: string): AIResponse {
+	private parseAIResponse(aiResponse: string, tablesInformation: Array<TableInformation>): AIResponse {
 		const cleanedResponse = cleanAIJsonResponse(aiResponse);
+		const tableNames = tablesInformation.map((t) => t.table_name);
 
 		try {
 			return JSON.parse(cleanedResponse) as AIResponse;
 		} catch (error) {
-			throw new Error(`Failed to parse AI response: ${error.message}`);
+			throw new Error(`Failed to parse AI response for tables [${tableNames.join(', ')}]: ${error.message}`);
 		}
 	}
 
@@ -156,7 +183,9 @@ Respond ONLY with valid JSON in this exact format (no markdown, no explanations)
 			settings.readonly_fields = this.filterValidColumns(tableSettings.readonly_fields, validColumnNames);
 			settings.columns_view = this.filterValidColumns(tableSettings.columns_view, validColumnNames);
 			settings.ordering = this.mapOrdering(tableSettings.ordering);
-			settings.ordering_field = validColumnNames.includes(tableSettings.ordering_field) ? tableSettings.ordering_field : null;
+			settings.ordering_field = validColumnNames.includes(tableSettings.ordering_field)
+				? tableSettings.ordering_field
+				: null;
 			settings.table_widgets = tableSettings.widgets
 				.filter((w) => validColumnNames.includes(w.field_name))
 				.map((widgetData) => {

backend/src/entities/shared-jobs/shared-jobs.service.ts

@@ -51,24 +51,44 @@ export class SharedJobsService {
 			}
 
 			const queue = new PQueue({ concurrency: 4 });
-			const tablesInformation = await Promise.all(
+			const tablesInformationResults = await Promise.all(
 				tablesToScan.map((table) =>
 					queue.add(async () => {
-						const structure = await dao.getTableStructure(table.tableName, null);
-						const primaryColumns = await dao.getTablePrimaryColumns(table.tableName, null);
-						const foreignKeys = await dao.getTableForeignKeys(table.tableName, null);
-						return {
-							table_name: table.tableName,
-							structure,
-							primaryColumns,
-							foreignKeys,
-						};
+						try {
+							const structure = await dao.getTableStructure(table.tableName, null);
+							const primaryColumns = await dao.getTablePrimaryColumns(table.tableName, null);
+							const foreignKeys = await dao.getTableForeignKeys(table.tableName, null);
+							return {
+								table_name: table.tableName,
+								structure,
+								primaryColumns,
+								foreignKeys,
+							};
+						} catch (error) {
+							console.error(`Error getting table information for "${table.tableName}": ${error.message}`);
+							return null;
+						}
 					}),
 				),
 			);
 
+			const tablesInformation = tablesInformationResults.filter((info) => info !== null);
+
+			if (tablesInformation.length === 0) {
+				console.info(`No valid tables to process for connection with id "${connection.id}"`);
+				return;
+			}
+
+			console.info(`Processing ${tablesInformation.length} tables with AI for connection "${connection.id}"`);
 			const generatedTableSettings = await this.aiService.generateNewTableSettingsWithAI(tablesInformation);
 
+			if (generatedTableSettings.length === 0) {
+				console.info(`No table settings generated by AI for connection with id "${connection.id}"`);
+				return;
+			}
+
+			console.info(`AI generated settings for ${generatedTableSettings.length} tables`);
+
 			const widgetsByTable = new Map<string, Array<TableWidgetEntity>>();
 			for (const setting of generatedTableSettings) {
 				if (setting.table_widgets && setting.table_widgets.length > 0) {