Merge pull request #1811 from rocket-admin/backend_ai_table_schema_reponses

feat: enhance AI schema change handling with clarification responses and improved DTOs

Author: lyubov-voloshko

backend/src/entities/table-schema/ai/run-schema-change-ai-loop.ts

@@ -34,10 +34,9 @@ export interface RunSchemaChangeAiLoopOptions {
 	logger?: Logger;
 }
 
-export interface SchemaChangeAiLoopResult {
-	proposals: ProposeSchemaChangeArgs[];
-	responseId: string | null;
-}
+export type SchemaChangeAiLoopResult =
+	| { kind: 'proposals'; proposals: ProposeSchemaChangeArgs[]; responseId: string | null }
+	| { kind: 'clarification'; assistantMessage: string; responseId: string | null };
 
 export async function runSchemaChangeAiLoop(opts: RunSchemaChangeAiLoopOptions): Promise<SchemaChangeAiLoopResult> {
 	const { aiCoreService, provider, dao, userEmail, logger } = opts;
@@ -77,15 +76,16 @@ export async function runSchemaChangeAiLoop(opts: RunSchemaChangeAiLoopOptions):
 		const proposalCall = pendingToolCalls.find((tc) => TERMINAL_PROPOSAL_TOOL_NAMES.has(tc.name));
 		if (proposalCall) {
 			const proposals = coerceAndValidateProposals(proposalCall);
-			return { proposals, responseId: lastResponseId };
+			return { kind: 'proposals', proposals, responseId: lastResponseId };
 		}
 
 		if (pendingToolCalls.length === 0) {
-			const hint = accumulatedContent
-				? `AI replied with text but no tool call: "${accumulatedContent.slice(0, 200)}"`
-				: 'AI produced no tool calls and no text.';
+			const trimmed = accumulatedContent.trim();
+			if (trimmed.length > 0) {
+				return { kind: 'clarification', assistantMessage: trimmed, responseId: lastResponseId };
+			}
 			throw new Error(
-				`${hint} The model must call proposeSchemaChange (SQL), proposeMongoSchemaChange (Mongo), proposeDynamoDbSchemaChange (DynamoDB), or proposeElasticsearchSchemaChange (Elasticsearch) with structured arguments.`,
+				'AI produced no tool calls and no text. The model must call proposeSchemaChange (SQL), proposeMongoSchemaChange (Mongo), proposeDynamoDbSchemaChange (DynamoDB), or proposeElasticsearchSchemaChange (Elasticsearch) with structured arguments, or ask a focused clarifying question in plain text.',
 			);
 		}
 

backend/src/entities/table-schema/ai/run-schema-change-fix-ai-loop.ts

@@ -0,0 +1,156 @@
+import { Logger } from '@nestjs/common';
+import { getDataAccessObject } from '@rocketadmin/shared-code/dist/src/data-access-layer/shared/create-data-access-object.js';
+import { ConnectionTypesEnum } from '@rocketadmin/shared-code/dist/src/shared/enums/connection-types-enum.js';
+import { AIProviderType } from '../../../ai-core/interfaces/ai-service.interface.js';
+import { AICoreService } from '../../../ai-core/services/ai-core.service.js';
+import { MessageBuilder } from '../../../ai-core/utils/message-builder.js';
+import { ConnectionEntity } from '../../connection/connection.entity.js';
+import { SchemaChangeTypeEnum } from '../table-schema-change-enums.js';
+import { isDynamoDbDialect, isElasticsearchDialect, isMongoDialect } from '../utils/assert-dialect-supported.js';
+import { runSchemaChangeAiLoop } from './run-schema-change-ai-loop.js';
+import {
+	createDynamoDbSchemaChangeTools,
+	createElasticsearchSchemaChangeTools,
+	createMongoSchemaChangeTools,
+	createSchemaChangeTools,
+	ProposeSchemaChangeArgs,
+} from './schema-change-tools.js';
+
+export interface SchemaChangeFixAiLoopOptions {
+	aiCoreService: AICoreService;
+	provider: AIProviderType;
+	connection: ConnectionEntity;
+	connectionType: ConnectionTypesEnum;
+	changeType: SchemaChangeTypeEnum;
+	targetTableName: string;
+	originalUserPrompt: string;
+	failingSql: string;
+	failingRollbackSql: string | null;
+	executionError: string;
+	logger?: Logger;
+}
+
+export interface SchemaChangeFixResult {
+	fixedProposal: ProposeSchemaChangeArgs;
+}
+
+export async function runSchemaChangeFixAiLoop(
+	opts: SchemaChangeFixAiLoopOptions,
+): Promise<SchemaChangeFixResult | null> {
+	const {
+		aiCoreService,
+		provider,
+		connection,
+		connectionType,
+		changeType,
+		targetTableName,
+		originalUserPrompt,
+		failingSql,
+		failingRollbackSql,
+		executionError,
+		logger,
+	} = opts;
+
+	const systemPrompt = buildFixSystemPrompt(connectionType, changeType, targetTableName);
+	const humanPrompt = buildFixHumanPrompt(originalUserPrompt, failingSql, failingRollbackSql, executionError);
+	const messages = new MessageBuilder().system(systemPrompt).human(humanPrompt).build();
+
+	const tools = isMongoDialect(connectionType)
+		? createMongoSchemaChangeTools()
+		: isDynamoDbDialect(connectionType)
+			? createDynamoDbSchemaChangeTools()
+			: isElasticsearchDialect(connectionType)
+				? createElasticsearchSchemaChangeTools()
+				: createSchemaChangeTools();
+
+	const dao = getDataAccessObject(connection);
+
+	const result = await runSchemaChangeAiLoop({
+		aiCoreService,
+		provider,
+		messages,
+		tools,
+		dao,
+		userEmail: undefined,
+		logger,
+		maxDepth: 4,
+	});
+
+	if (result.kind !== 'proposals') {
+		logger?.warn(`AI fix loop returned ${result.kind} instead of proposals; treating as fix-failed.`);
+		return null;
+	}
+
+	if (result.proposals.length !== 1) {
+		logger?.warn(
+			`AI fix loop returned ${result.proposals.length} proposals; expected exactly 1. Treating as fix-failed.`,
+		);
+		return null;
+	}
+
+	const proposal = result.proposals[0];
+
+	if (proposal.changeType !== changeType) {
+		logger?.warn(
+			`AI fix loop changed changeType from ${changeType} to ${proposal.changeType}; treating as fix-failed.`,
+		);
+		return null;
+	}
+	if (proposal.targetTableName !== targetTableName) {
+		logger?.warn(
+			`AI fix loop changed targetTableName from ${targetTableName} to ${proposal.targetTableName}; treating as fix-failed.`,
+		);
+		return null;
+	}
+	if (!proposal.forwardSql || proposal.forwardSql.trim() === failingSql.trim()) {
+		logger?.warn('AI fix loop returned an identical or empty forwardSql; treating as fix-failed.');
+		return null;
+	}
+
+	return { fixedProposal: proposal };
+}
+
+function buildFixSystemPrompt(
+	connectionType: ConnectionTypesEnum,
+	changeType: SchemaChangeTypeEnum,
+	targetTableName: string,
+): string {
+	return `You are a DDL repair assistant for ${connectionType}.
+
+A previously proposed schema change failed when applied against the live database. Your single job is to repair the failing statement and emit a corrected proposal via the appropriate proposeSchemaChange tool.
+
+Hard constraints (violations will be rejected and the fix discarded):
+- Emit EXACTLY ONE proposal.
+- The proposal MUST keep changeType = "${changeType}".
+- The proposal MUST keep targetTableName = "${targetTableName}".
+- The corrected forwardSql MUST be different from the failing forwardSql.
+- All previously documented rules for ${changeType} on ${connectionType} still apply (single statement, dialect-correct identifier quoting, no DML, etc.).
+
+Workflow:
+1. Read the database error carefully. Identify the precise cause (missing dependency, wrong type, reserved word, conflicting name, etc.).
+2. If — and only if — the fix depends on inspecting the live structure of an existing table or collection, call getTableStructure first.
+3. Emit the corrected proposal via the appropriate proposeSchemaChange / proposeMongoSchemaChange / proposeDynamoDbSchemaChange / proposeElasticsearchSchemaChange tool.
+
+Do NOT reply with free text. If you cannot produce a corrected statement, still emit a proposal with your best attempt and explain the residual risk in "reasoning".`;
+}
+
+function buildFixHumanPrompt(
+	originalUserPrompt: string,
+	failingSql: string,
+	failingRollbackSql: string | null,
+	executionError: string,
+): string {
+	return `Original user request:
+${originalUserPrompt}
+
+Failing forwardSql:
+${failingSql}
+
+Previously proposed rollbackSql:
+${failingRollbackSql ?? '(none)'}
+
+Database error returned when applying forwardSql:
+${executionError}
+
+Produce a corrected proposal that addresses the database error while still fulfilling the user's original request.`;
+}

backend/src/entities/table-schema/ai/schema-change-prompts.ts

@@ -37,6 +37,7 @@ Your job: translate the user's natural-language request into ONE OR MORE DDL sta
 Workflow:
 1. If the user's request references an existing table, call getTableStructure FIRST to inspect it. You may call it multiple times for different tables.
 2. Once you have enough context, call proposeSchemaChange EXACTLY ONCE with a "proposals" array containing every change required by the user's request. Do not write free-text explanations outside the tool call.
+3. If — and ONLY if — the user's request is genuinely too vague to translate into ANY concrete schema change (e.g. "draw a database", "make me something", with no domain, tables, or fields hinted), reply with a single short plain-text message containing ONE focused clarifying question (no tool call). Do not invent a domain. Once the user answers, proceed to step 2.
 
 Multi-proposal rules:
 - Order proposals so that any change depending on another comes AFTER its dependency. Tables referenced by foreign keys must be created BEFORE the tables that reference them.
@@ -85,6 +86,7 @@ DynamoDB is a NoSQL key-value store with a fixed primary key schema per table an
 Workflow:
 1. If the user's request references an existing table, call getTableStructure FIRST to see its key schema and sampled attributes. You may call it multiple times.
 2. Once you have enough context, call proposeDynamoDbSchemaChange EXACTLY ONCE with a "proposals" array containing every change required by the user's request. Do not write free-text explanations outside the tool call.
+3. If — and ONLY if — the user's request is genuinely too vague to translate into ANY concrete schema change (e.g. "draw a database", "make me something", with no domain, tables, or attributes hinted), reply with a single short plain-text message containing ONE focused clarifying question (no tool call). Do not invent a domain. Once the user answers, proceed to step 2.
 
 Multi-proposal rules:
 - Order proposals so dependent changes come after their prerequisites.
@@ -168,6 +170,7 @@ Elasticsearch is a search engine where each "table" is an index with a mapping (
 Workflow:
 1. If the user's request references an existing index, call getTableStructure FIRST to inspect its sampled fields and inferred types. You may call it multiple times.
 2. Once you have enough context, call proposeElasticsearchSchemaChange EXACTLY ONCE with a "proposals" array containing every change required by the user's request. Do not write free-text explanations outside the tool call.
+3. If — and ONLY if — the user's request is genuinely too vague to translate into ANY concrete schema change (e.g. "draw a database", "make me something", with no domain, indices, or fields hinted), reply with a single short plain-text message containing ONE focused clarifying question (no tool call). Do not invent a domain. Once the user answers, proceed to step 2.
 
 Multi-proposal rules:
 - Order proposals so dependent changes come after their prerequisites.
@@ -225,6 +228,7 @@ MongoDB does not use SQL DDL. Instead, schema changes are structured JSON operat
 Workflow:
 1. If the user's request references an existing collection, call getTableStructure FIRST to inspect a sample of documents. You may call it multiple times for different collections.
 2. Once you have enough context, call proposeMongoSchemaChange EXACTLY ONCE with a "proposals" array containing every change required by the user's request. Do not write free-text explanations outside the tool call.
+3. If — and ONLY if — the user's request is genuinely too vague to translate into ANY concrete schema change (e.g. "draw a database", "make me something", with no domain, collections, or fields hinted), reply with a single short plain-text message containing ONE focused clarifying question (no tool call). Do not invent a domain. Once the user answers, proceed to step 2.
 
 Multi-proposal rules:
 - Order proposals so dependent changes come after their prerequisites (e.g. createCollection before createIndex on that collection).

backend/src/entities/table-schema/application/data-transfer-objects/schema-change-batch-response.dto.ts

@@ -3,14 +3,17 @@ import { SchemaChangeResponseDto } from './schema-change-response.dto.js';
 
 export class SchemaChangeBatchResponseDto {
 	@ApiProperty({
+		required: false,
+		nullable: true,
 		description:
-			'Identifier shared by every change generated from a single user prompt. Use it to approve/reject/rollback the entire batch in one call.',
+			'Identifier shared by every change generated from a single user prompt. Use it to approve/reject/rollback the entire batch in one call. Null when the AI did not propose any change (e.g. it returned a clarifying question instead — see assistantMessage).',
 	})
-	batchId: string;
+	batchId: string | null;
 
 	@ApiProperty({
 		type: [SchemaChangeResponseDto],
-		description: 'Generated changes ordered by orderInBatch (dependency order — parents first).',
+		description:
+			'Generated changes ordered by orderInBatch (dependency order — parents first). Empty when the AI returned a clarifying question instead of a proposal.',
 	})
 	changes: SchemaChangeResponseDto[];
 
@@ -21,4 +24,12 @@ export class SchemaChangeBatchResponseDto {
 			'Conversation thread ID. Present when the request used or created a chat thread. Pass it back as the threadId query param on subsequent generate calls to continue the conversation with full prior context.',
 	})
 	threadId?: string | null;
+
+	@ApiProperty({
+		required: false,
+		nullable: true,
+		description:
+			"Free-text reply from the AI when it could not propose any change and needs more information from the user (e.g. a clarifying question). When present, batchId is null and changes is empty; resubmit the user's answer with the same threadId to continue the conversation.",
+	})
+	assistantMessage?: string | null;
 }

backend/src/entities/table-schema/application/data-transfer-objects/schema-change-response.dto.ts

@@ -54,6 +54,25 @@ export class SchemaChangeResponseDto {
 	@ApiProperty()
 	autoRollbackSucceeded: boolean;
 
+	@ApiProperty({
+		description:
+			'True when the original AI-proposed SQL failed at apply time and the AI was asked to repair it. When true, forwardSql/rollbackSql reflect the repaired statements and the originals are preserved in aiAutoFixOriginalForwardSql/aiAutoFixOriginalRollbackSql.',
+	})
+	aiAutoFixApplied: boolean;
+
+	@ApiProperty({ required: false, nullable: true })
+	aiAutoFixOriginalForwardSql: string | null;
+
+	@ApiProperty({ required: false, nullable: true })
+	aiAutoFixOriginalRollbackSql: string | null;
+
+	@ApiProperty({
+		required: false,
+		nullable: true,
+		description: 'Database error returned when the original SQL was attempted, before the AI repaired it.',
+	})
+	aiAutoFixOriginalError: string | null;
+
 	@ApiProperty()
 	userPrompt: string;
 

backend/src/entities/table-schema/repository/custom-table-schema-change-repository-extension.ts

@@ -77,6 +77,24 @@ export const customTableSchemaChangeRepositoryExtension: ITableSchemaChangeRepos
 			if (meta.userModifiedSql !== undefined) {
 				patch.userModifiedSql = meta.userModifiedSql;
 			}
+			if (meta.forwardSql !== undefined) {
+				patch.forwardSql = meta.forwardSql;
+			}
+			if (meta.rollbackSql !== undefined) {
+				patch.rollbackSql = meta.rollbackSql;
+			}
+			if (meta.aiAutoFixApplied !== undefined) {
+				patch.aiAutoFixApplied = meta.aiAutoFixApplied;
+			}
+			if (meta.aiAutoFixOriginalForwardSql !== undefined) {
+				patch.aiAutoFixOriginalForwardSql = meta.aiAutoFixOriginalForwardSql;
+			}
+			if (meta.aiAutoFixOriginalRollbackSql !== undefined) {
+				patch.aiAutoFixOriginalRollbackSql = meta.aiAutoFixOriginalRollbackSql;
+			}
+			if (meta.aiAutoFixOriginalError !== undefined) {
+				patch.aiAutoFixOriginalError = meta.aiAutoFixOriginalError;
+			}
 		}
 		await this.update({ id }, patch);
 		return await this.findOne({ where: { id } });

backend/src/entities/table-schema/repository/table-schema-change.repository.interface.ts

@@ -9,6 +9,12 @@ export interface UpdateSchemaChangeStatusMeta {
 	autoRollbackAttempted?: boolean;
 	autoRollbackSucceeded?: boolean;
 	userModifiedSql?: string | null;
+	forwardSql?: string;
+	rollbackSql?: string | null;
+	aiAutoFixApplied?: boolean;
+	aiAutoFixOriginalForwardSql?: string | null;
+	aiAutoFixOriginalRollbackSql?: string | null;
+	aiAutoFixOriginalError?: string | null;
 }
 
 export interface ITableSchemaChangeRepository {

backend/src/entities/table-schema/table-schema-change.entity.ts

@@ -88,6 +88,18 @@ export class TableSchemaChangeEntity {
 	@Column({ type: 'varchar', length: 128, nullable: true })
 	aiModelUsed: string | null;
 
+	@Column({ type: 'boolean', default: false })
+	aiAutoFixApplied: boolean;
+
+	@Column({ type: 'text', nullable: true })
+	aiAutoFixOriginalForwardSql: string | null;
+
+	@Column({ type: 'text', nullable: true })
+	aiAutoFixOriginalRollbackSql: string | null;
+
+	@Column({ type: 'text', nullable: true })
+	aiAutoFixOriginalError: string | null;
+
 	@Column({ type: 'timestamp', default: () => 'CURRENT_TIMESTAMP' })
 	createdAt: Date;
 

backend/src/entities/table-schema/table-schema.controller.ts

@@ -72,7 +72,7 @@ export class TableSchemaController {
 
 	@ApiOperation({
 		summary:
-			'Generate one or more schema changes from a natural-language prompt. The response is always a batch envelope; single-change prompts return a length-1 array. Pass an optional threadId in the body to continue an existing conversation; the response returns the threadId to use for follow-up turns.',
+			"Generate one or more schema changes from a natural-language prompt. The response is always a batch envelope; single-change prompts return a length-1 array. Pass an optional threadId in the body to continue an existing conversation; the response returns the threadId to use for follow-up turns. If the AI needs more context (e.g. the request is too vague), the response carries an assistantMessage with a clarifying question, batchId is null, and changes is empty — resubmit the user's answer with the same threadId to continue.",
 	})
 	@ApiParam({ name: 'connectionId', type: String })
 	@ApiBody({ type: GenerateSchemaChangeDto })