Skip to content

Commit df1adbf

Browse files
ArtuomkaArtuomka
committed
feat: implement AI auto-fix functionality for schema changes with new database fields and response DTOs
1 parent 168d21d commit df1adbf

11 files changed

Lines changed: 605 additions & 14 deletions

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

Lines changed: 156 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,156 @@
1+
import { Logger } from '@nestjs/common';
2+
import { getDataAccessObject } from '@rocketadmin/shared-code/dist/src/data-access-layer/shared/create-data-access-object.js';
3+
import { ConnectionTypesEnum } from '@rocketadmin/shared-code/dist/src/shared/enums/connection-types-enum.js';
4+
import { AIProviderType } from '../../../ai-core/interfaces/ai-service.interface.js';
5+
import { AICoreService } from '../../../ai-core/services/ai-core.service.js';
6+
import { MessageBuilder } from '../../../ai-core/utils/message-builder.js';
7+
import { ConnectionEntity } from '../../connection/connection.entity.js';
8+
import { SchemaChangeTypeEnum } from '../table-schema-change-enums.js';
9+
import { isDynamoDbDialect, isElasticsearchDialect, isMongoDialect } from '../utils/assert-dialect-supported.js';
10+
import { runSchemaChangeAiLoop } from './run-schema-change-ai-loop.js';
11+
import {
12+
createDynamoDbSchemaChangeTools,
13+
createElasticsearchSchemaChangeTools,
14+
createMongoSchemaChangeTools,
15+
createSchemaChangeTools,
16+
ProposeSchemaChangeArgs,
17+
} from './schema-change-tools.js';
18+
19+
export interface SchemaChangeFixAiLoopOptions {
20+
aiCoreService: AICoreService;
21+
provider: AIProviderType;
22+
connection: ConnectionEntity;
23+
connectionType: ConnectionTypesEnum;
24+
changeType: SchemaChangeTypeEnum;
25+
targetTableName: string;
26+
originalUserPrompt: string;
27+
failingSql: string;
28+
failingRollbackSql: string | null;
29+
executionError: string;
30+
logger?: Logger;
31+
}
32+
33+
export interface SchemaChangeFixResult {
34+
fixedProposal: ProposeSchemaChangeArgs;
35+
}
36+
37+
export async function runSchemaChangeFixAiLoop(
38+
opts: SchemaChangeFixAiLoopOptions,
39+
): Promise<SchemaChangeFixResult | null> {
40+
const {
41+
aiCoreService,
42+
provider,
43+
connection,
44+
connectionType,
45+
changeType,
46+
targetTableName,
47+
originalUserPrompt,
48+
failingSql,
49+
failingRollbackSql,
50+
executionError,
51+
logger,
52+
} = opts;
53+
54+
const systemPrompt = buildFixSystemPrompt(connectionType, changeType, targetTableName);
55+
const humanPrompt = buildFixHumanPrompt(originalUserPrompt, failingSql, failingRollbackSql, executionError);
56+
const messages = new MessageBuilder().system(systemPrompt).human(humanPrompt).build();
57+
58+
const tools = isMongoDialect(connectionType)
59+
? createMongoSchemaChangeTools()
60+
: isDynamoDbDialect(connectionType)
61+
? createDynamoDbSchemaChangeTools()
62+
: isElasticsearchDialect(connectionType)
63+
? createElasticsearchSchemaChangeTools()
64+
: createSchemaChangeTools();
65+
66+
const dao = getDataAccessObject(connection);
67+
68+
const result = await runSchemaChangeAiLoop({
69+
aiCoreService,
70+
provider,
71+
messages,
72+
tools,
73+
dao,
74+
userEmail: undefined,
75+
logger,
76+
maxDepth: 4,
77+
});
78+
79+
if (result.kind !== 'proposals') {
80+
logger?.warn(`AI fix loop returned ${result.kind} instead of proposals; treating as fix-failed.`);
81+
return null;
82+
}
83+
84+
if (result.proposals.length !== 1) {
85+
logger?.warn(
86+
`AI fix loop returned ${result.proposals.length} proposals; expected exactly 1. Treating as fix-failed.`,
87+
);
88+
return null;
89+
}
90+
91+
const proposal = result.proposals[0];
92+
93+
if (proposal.changeType !== changeType) {
94+
logger?.warn(
95+
`AI fix loop changed changeType from ${changeType} to ${proposal.changeType}; treating as fix-failed.`,
96+
);
97+
return null;
98+
}
99+
if (proposal.targetTableName !== targetTableName) {
100+
logger?.warn(
101+
`AI fix loop changed targetTableName from ${targetTableName} to ${proposal.targetTableName}; treating as fix-failed.`,
102+
);
103+
return null;
104+
}
105+
if (!proposal.forwardSql || proposal.forwardSql.trim() === failingSql.trim()) {
106+
logger?.warn('AI fix loop returned an identical or empty forwardSql; treating as fix-failed.');
107+
return null;
108+
}
109+
110+
return { fixedProposal: proposal };
111+
}
112+
113+
function buildFixSystemPrompt(
114+
connectionType: ConnectionTypesEnum,
115+
changeType: SchemaChangeTypeEnum,
116+
targetTableName: string,
117+
): string {
118+
return `You are a DDL repair assistant for ${connectionType}.
119+
120+
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.
121+
122+
Hard constraints (violations will be rejected and the fix discarded):
123+
- Emit EXACTLY ONE proposal.
124+
- The proposal MUST keep changeType = "${changeType}".
125+
- The proposal MUST keep targetTableName = "${targetTableName}".
126+
- The corrected forwardSql MUST be different from the failing forwardSql.
127+
- All previously documented rules for ${changeType} on ${connectionType} still apply (single statement, dialect-correct identifier quoting, no DML, etc.).
128+
129+
Workflow:
130+
1. Read the database error carefully. Identify the precise cause (missing dependency, wrong type, reserved word, conflicting name, etc.).
131+
2. If — and only if — the fix depends on inspecting the live structure of an existing table or collection, call getTableStructure first.
132+
3. Emit the corrected proposal via the appropriate proposeSchemaChange / proposeMongoSchemaChange / proposeDynamoDbSchemaChange / proposeElasticsearchSchemaChange tool.
133+
134+
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".`;
135+
}
136+
137+
function buildFixHumanPrompt(
138+
originalUserPrompt: string,
139+
failingSql: string,
140+
failingRollbackSql: string | null,
141+
executionError: string,
142+
): string {
143+
return `Original user request:
144+
${originalUserPrompt}
145+
146+
Failing forwardSql:
147+
${failingSql}
148+
149+
Previously proposed rollbackSql:
150+
${failingRollbackSql ?? '(none)'}
151+
152+
Database error returned when applying forwardSql:
153+
${executionError}
154+
155+
Produce a corrected proposal that addresses the database error while still fulfilling the user's original request.`;
156+
}

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

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -54,6 +54,25 @@ export class SchemaChangeResponseDto {
5454
@ApiProperty()
5555
autoRollbackSucceeded: boolean;
5656

57+
@ApiProperty({
58+
description:
59+
'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.',
60+
})
61+
aiAutoFixApplied: boolean;
62+
63+
@ApiProperty({ required: false, nullable: true })
64+
aiAutoFixOriginalForwardSql: string | null;
65+
66+
@ApiProperty({ required: false, nullable: true })
67+
aiAutoFixOriginalRollbackSql: string | null;
68+
69+
@ApiProperty({
70+
required: false,
71+
nullable: true,
72+
description: 'Database error returned when the original SQL was attempted, before the AI repaired it.',
73+
})
74+
aiAutoFixOriginalError: string | null;
75+
5776
@ApiProperty()
5877
userPrompt: string;
5978

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

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -77,6 +77,24 @@ export const customTableSchemaChangeRepositoryExtension: ITableSchemaChangeRepos
7777
if (meta.userModifiedSql !== undefined) {
7878
patch.userModifiedSql = meta.userModifiedSql;
7979
}
80+
if (meta.forwardSql !== undefined) {
81+
patch.forwardSql = meta.forwardSql;
82+
}
83+
if (meta.rollbackSql !== undefined) {
84+
patch.rollbackSql = meta.rollbackSql;
85+
}
86+
if (meta.aiAutoFixApplied !== undefined) {
87+
patch.aiAutoFixApplied = meta.aiAutoFixApplied;
88+
}
89+
if (meta.aiAutoFixOriginalForwardSql !== undefined) {
90+
patch.aiAutoFixOriginalForwardSql = meta.aiAutoFixOriginalForwardSql;
91+
}
92+
if (meta.aiAutoFixOriginalRollbackSql !== undefined) {
93+
patch.aiAutoFixOriginalRollbackSql = meta.aiAutoFixOriginalRollbackSql;
94+
}
95+
if (meta.aiAutoFixOriginalError !== undefined) {
96+
patch.aiAutoFixOriginalError = meta.aiAutoFixOriginalError;
97+
}
8098
}
8199
await this.update({ id }, patch);
82100
return await this.findOne({ where: { id } });

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

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,12 @@ export interface UpdateSchemaChangeStatusMeta {
99
autoRollbackAttempted?: boolean;
1010
autoRollbackSucceeded?: boolean;
1111
userModifiedSql?: string | null;
12+
forwardSql?: string;
13+
rollbackSql?: string | null;
14+
aiAutoFixApplied?: boolean;
15+
aiAutoFixOriginalForwardSql?: string | null;
16+
aiAutoFixOriginalRollbackSql?: string | null;
17+
aiAutoFixOriginalError?: string | null;
1218
}
1319

1420
export interface ITableSchemaChangeRepository {

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

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -88,6 +88,18 @@ export class TableSchemaChangeEntity {
8888
@Column({ type: 'varchar', length: 128, nullable: true })
8989
aiModelUsed: string | null;
9090

91+
@Column({ type: 'boolean', default: false })
92+
aiAutoFixApplied: boolean;
93+
94+
@Column({ type: 'text', nullable: true })
95+
aiAutoFixOriginalForwardSql: string | null;
96+
97+
@Column({ type: 'text', nullable: true })
98+
aiAutoFixOriginalRollbackSql: string | null;
99+
100+
@Column({ type: 'text', nullable: true })
101+
aiAutoFixOriginalError: string | null;
102+
91103
@Column({ type: 'timestamp', default: () => 'CURRENT_TIMESTAMP' })
92104
createdAt: Date;
93105

0 commit comments

Comments
 (0)