diff --git a/backend/src/common/application/global-database-context.interface.ts b/backend/src/common/application/global-database-context.interface.ts index 9ee2ef771..2dc8616b1 100644 --- a/backend/src/common/application/global-database-context.interface.ts +++ b/backend/src/common/application/global-database-context.interface.ts @@ -37,6 +37,10 @@ import { TableFiltersEntity } from '../../entities/table-filters/table-filters.e import { TableInfoEntity } from '../../entities/table-info/table-info.entity.js'; import { ITableLogsRepository } from '../../entities/table-logs/repository/table-logs-repository.interface.js'; import { ITableSchemaChangeRepository } from '../../entities/table-schema/repository/table-schema-change.repository.interface.js'; +import { ISchemaChangeChatRepository } from '../../entities/table-schema/schema-change-chat/schema-change-chat/repository/schema-change-chat-repository.interface.js'; +import { SchemaChangeChatEntity } from '../../entities/table-schema/schema-change-chat/schema-change-chat/schema-change-chat.entity.js'; +import { ISchemaChangeChatMessageRepository } from '../../entities/table-schema/schema-change-chat/schema-change-chat-message/repository/schema-change-chat-message-repository.interface.js'; +import { SchemaChangeChatMessageEntity } from '../../entities/table-schema/schema-change-chat/schema-change-chat-message/schema-change-chat-message.entity.js'; import { TableSchemaChangeEntity } from '../../entities/table-schema/table-schema-change.entity.js'; import { ITableSettingsRepository } from '../../entities/table-settings/common-table-settings/repository/table-settings.repository.interface.js'; import { TableSettingsEntity } from '../../entities/table-settings/common-table-settings/table-settings.entity.js'; @@ -106,4 +110,6 @@ export interface IGlobalDatabaseContext extends IDatabaseContext { userAiChatRepository: Repository & IUserAiChatRepository; aiChatMessageRepository: Repository & IAiChatMessageRepository; tableSchemaChangeRepository: Repository & ITableSchemaChangeRepository; + schemaChangeChatRepository: Repository & ISchemaChangeChatRepository; + schemaChangeChatMessageRepository: Repository & ISchemaChangeChatMessageRepository; } diff --git a/backend/src/common/application/global-database-context.ts b/backend/src/common/application/global-database-context.ts index 43d7e5e31..0a300738b 100644 --- a/backend/src/common/application/global-database-context.ts +++ b/backend/src/common/application/global-database-context.ts @@ -64,6 +64,12 @@ import { ITableLogsRepository } from '../../entities/table-logs/repository/table import { TableLogsEntity } from '../../entities/table-logs/table-logs.entity.js'; import { customTableSchemaChangeRepositoryExtension } from '../../entities/table-schema/repository/custom-table-schema-change-repository-extension.js'; import { ITableSchemaChangeRepository } from '../../entities/table-schema/repository/table-schema-change.repository.interface.js'; +import { schemaChangeChatRepositoryExtension } from '../../entities/table-schema/schema-change-chat/schema-change-chat/repository/schema-change-chat-repository.extension.js'; +import { ISchemaChangeChatRepository } from '../../entities/table-schema/schema-change-chat/schema-change-chat/repository/schema-change-chat-repository.interface.js'; +import { SchemaChangeChatEntity } from '../../entities/table-schema/schema-change-chat/schema-change-chat/schema-change-chat.entity.js'; +import { schemaChangeChatMessageRepositoryExtension } from '../../entities/table-schema/schema-change-chat/schema-change-chat-message/repository/schema-change-chat-message-repository.extension.js'; +import { ISchemaChangeChatMessageRepository } from '../../entities/table-schema/schema-change-chat/schema-change-chat-message/repository/schema-change-chat-message-repository.interface.js'; +import { SchemaChangeChatMessageEntity } from '../../entities/table-schema/schema-change-chat/schema-change-chat-message/schema-change-chat-message.entity.js'; import { TableSchemaChangeEntity } from '../../entities/table-schema/table-schema-change.entity.js'; import { ITableSettingsRepository } from '../../entities/table-settings/common-table-settings/repository/table-settings.repository.interface.js'; import { tableSettingsCustomRepositoryExtension } from '../../entities/table-settings/common-table-settings/repository/table-settings-custom-repository-extension.js'; @@ -157,6 +163,9 @@ export class GlobalDatabaseContext implements IGlobalDatabaseContext { private _userAiChatRepository: Repository & IUserAiChatRepository; private _aiChatMessageRepository: Repository & IAiChatMessageRepository; private _tableSchemaChangeRepository: Repository & ITableSchemaChangeRepository; + private _schemaChangeChatRepository: Repository & ISchemaChangeChatRepository; + private _schemaChangeChatMessageRepository: Repository & + ISchemaChangeChatMessageRepository; public constructor( @Inject(BaseType.DATA_SOURCE) @@ -264,6 +273,12 @@ export class GlobalDatabaseContext implements IGlobalDatabaseContext { this._tableSchemaChangeRepository = this.appDataSource .getRepository(TableSchemaChangeEntity) .extend(customTableSchemaChangeRepositoryExtension); + this._schemaChangeChatRepository = this.appDataSource + .getRepository(SchemaChangeChatEntity) + .extend(schemaChangeChatRepositoryExtension); + this._schemaChangeChatMessageRepository = this.appDataSource + .getRepository(SchemaChangeChatMessageEntity) + .extend(schemaChangeChatMessageRepositoryExtension); } public get userRepository(): Repository & IUserRepository { @@ -428,6 +443,15 @@ export class GlobalDatabaseContext implements IGlobalDatabaseContext { return this._tableSchemaChangeRepository; } + public get schemaChangeChatRepository(): Repository & ISchemaChangeChatRepository { + return this._schemaChangeChatRepository; + } + + public get schemaChangeChatMessageRepository(): Repository & + ISchemaChangeChatMessageRepository { + return this._schemaChangeChatMessageRepository; + } + public startTransaction(): Promise { this._queryRunner = this.appDataSource.createQueryRunner(); this._queryRunner.startTransaction(); diff --git a/backend/src/entities/table-schema/application/data-structures/generate-schema-change.ds.ts b/backend/src/entities/table-schema/application/data-structures/generate-schema-change.ds.ts index a73504841..233ce7f48 100644 --- a/backend/src/entities/table-schema/application/data-structures/generate-schema-change.ds.ts +++ b/backend/src/entities/table-schema/application/data-structures/generate-schema-change.ds.ts @@ -3,4 +3,5 @@ export class GenerateSchemaChangeDs { userPrompt: string; userId: string; masterPassword?: string; + threadId?: string | null; } diff --git a/backend/src/entities/table-schema/application/data-transfer-objects/generate-schema-change.dto.ts b/backend/src/entities/table-schema/application/data-transfer-objects/generate-schema-change.dto.ts index c7d9ed7f2..72e45933d 100644 --- a/backend/src/entities/table-schema/application/data-transfer-objects/generate-schema-change.dto.ts +++ b/backend/src/entities/table-schema/application/data-transfer-objects/generate-schema-change.dto.ts @@ -1,5 +1,5 @@ import { ApiProperty } from '@nestjs/swagger'; -import { IsNotEmpty, IsString, MaxLength, MinLength } from 'class-validator'; +import { IsNotEmpty, IsOptional, IsString, IsUUID, MaxLength, MinLength } from 'class-validator'; export class GenerateSchemaChangeDto { @ApiProperty({ @@ -15,4 +15,15 @@ export class GenerateSchemaChangeDto { @MinLength(1) @MaxLength(2000) userPrompt: string; + + @ApiProperty({ + type: String, + required: false, + nullable: true, + description: + 'Optional thread ID to continue an existing conversation. When supplied, prior turns are prepended to the AI prompt, giving the model context for iterative refinement (e.g. "now also add an index", "rename it to created_at"). Omit to start a fresh thread; the returned threadId can be passed back on the next call.', + }) + @IsOptional() + @IsUUID() + threadId?: string; } diff --git a/backend/src/entities/table-schema/application/data-transfer-objects/schema-change-batch-response.dto.ts b/backend/src/entities/table-schema/application/data-transfer-objects/schema-change-batch-response.dto.ts index f0c9281fb..b4ae87970 100644 --- a/backend/src/entities/table-schema/application/data-transfer-objects/schema-change-batch-response.dto.ts +++ b/backend/src/entities/table-schema/application/data-transfer-objects/schema-change-batch-response.dto.ts @@ -13,4 +13,12 @@ export class SchemaChangeBatchResponseDto { description: 'Generated changes ordered by orderInBatch (dependency order — parents first).', }) changes: SchemaChangeResponseDto[]; + + @ApiProperty({ + required: false, + nullable: true, + description: + '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; } diff --git a/backend/src/entities/table-schema/schema-change-chat/schema-change-chat-message/repository/schema-change-chat-message-repository.extension.ts b/backend/src/entities/table-schema/schema-change-chat/schema-change-chat-message/repository/schema-change-chat-message-repository.extension.ts new file mode 100644 index 000000000..6e52d3a09 --- /dev/null +++ b/backend/src/entities/table-schema/schema-change-chat/schema-change-chat-message/repository/schema-change-chat-message-repository.extension.ts @@ -0,0 +1,35 @@ +import { MessageRole } from '../../../../ai/ai-conversation-history/ai-chat-messages/message-role.enum.js'; +import { SchemaChangeChatMessageEntity } from '../schema-change-chat-message.entity.js'; +import { ISchemaChangeChatMessageRepository } from './schema-change-chat-message-repository.interface.js'; + +export const schemaChangeChatMessageRepositoryExtension: ISchemaChangeChatMessageRepository = { + async findMessagesForChat(chatId: string): Promise { + return await this.createQueryBuilder('schema_change_chat_message') + .where('schema_change_chat_message.chat_id = :chatId', { chatId }) + .orderBy('schema_change_chat_message.created_at', 'ASC') + .getMany(); + }, + + async deleteMessagesForChat(chatId: string): Promise { + await this.createQueryBuilder() + .delete() + .from('schema_change_chat_message') + .where('chat_id = :chatId', { chatId }) + .execute(); + }, + + async saveMessage( + chatId: string, + message: string, + role: MessageRole, + batchId?: string | null, + ): Promise { + const newMessage = this.create({ + chat_id: chatId, + message, + role, + batch_id: batchId ?? null, + }); + return await this.save(newMessage); + }, +}; diff --git a/backend/src/entities/table-schema/schema-change-chat/schema-change-chat-message/repository/schema-change-chat-message-repository.interface.ts b/backend/src/entities/table-schema/schema-change-chat/schema-change-chat-message/repository/schema-change-chat-message-repository.interface.ts new file mode 100644 index 000000000..5c0e232fc --- /dev/null +++ b/backend/src/entities/table-schema/schema-change-chat/schema-change-chat-message/repository/schema-change-chat-message-repository.interface.ts @@ -0,0 +1,13 @@ +import { MessageRole } from '../../../../ai/ai-conversation-history/ai-chat-messages/message-role.enum.js'; +import { SchemaChangeChatMessageEntity } from '../schema-change-chat-message.entity.js'; + +export interface ISchemaChangeChatMessageRepository { + findMessagesForChat(chatId: string): Promise; + deleteMessagesForChat(chatId: string): Promise; + saveMessage( + chatId: string, + message: string, + role: MessageRole, + batchId?: string | null, + ): Promise; +} diff --git a/backend/src/entities/table-schema/schema-change-chat/schema-change-chat-message/schema-change-chat-message.entity.ts b/backend/src/entities/table-schema/schema-change-chat/schema-change-chat-message/schema-change-chat-message.entity.ts new file mode 100644 index 000000000..f5a8dd4ae --- /dev/null +++ b/backend/src/entities/table-schema/schema-change-chat/schema-change-chat-message/schema-change-chat-message.entity.ts @@ -0,0 +1,44 @@ +import { + Column, + CreateDateColumn, + Entity, + JoinColumn, + ManyToOne, + PrimaryGeneratedColumn, + Relation, + UpdateDateColumn, +} from 'typeorm'; +import { MessageRole } from '../../../ai/ai-conversation-history/ai-chat-messages/message-role.enum.js'; +import { SchemaChangeChatEntity } from '../schema-change-chat/schema-change-chat.entity.js'; + +@Entity('schema_change_chat_message') +export class SchemaChangeChatMessageEntity { + @PrimaryGeneratedColumn('uuid') + id: string; + + @Column({ default: null, type: 'text' }) + message: string; + + @Column({ nullable: true, default: null, type: 'enum', enum: MessageRole }) + role: MessageRole; + + @Column({ type: 'uuid', nullable: true, default: null }) + batch_id: string | null; + + @CreateDateColumn({ type: 'timestamp', default: () => 'CURRENT_TIMESTAMP' }) + created_at: Date; + + @UpdateDateColumn({ type: 'timestamp', nullable: true, default: null }) + updated_at: Date; + + @ManyToOne( + () => SchemaChangeChatEntity, + (chat) => chat.messages, + { onDelete: 'CASCADE' }, + ) + @JoinColumn({ name: 'chat_id' }) + chat: Relation; + + @Column() + chat_id: string; +} diff --git a/backend/src/entities/table-schema/schema-change-chat/schema-change-chat/repository/schema-change-chat-repository.extension.ts b/backend/src/entities/table-schema/schema-change-chat/schema-change-chat/repository/schema-change-chat-repository.extension.ts new file mode 100644 index 000000000..01ea497d7 --- /dev/null +++ b/backend/src/entities/table-schema/schema-change-chat/schema-change-chat/repository/schema-change-chat-repository.extension.ts @@ -0,0 +1,53 @@ +import { SchemaChangeChatEntity } from '../schema-change-chat.entity.js'; +import { ISchemaChangeChatRepository } from './schema-change-chat-repository.interface.js'; + +export const schemaChangeChatRepositoryExtension: ISchemaChangeChatRepository = { + async findChatByIdAndUserId(chatId: string, userId: string): Promise { + return await this.createQueryBuilder('schema_change_chat') + .where('schema_change_chat.id = :chatId', { chatId }) + .andWhere('schema_change_chat.user_id = :userId', { userId }) + .getOne(); + }, + + async findChatWithMessagesByIdAndUserId(chatId: string, userId: string): Promise { + return await this.createQueryBuilder('schema_change_chat') + .leftJoinAndSelect('schema_change_chat.messages', 'messages') + .where('schema_change_chat.id = :chatId', { chatId }) + .andWhere('schema_change_chat.user_id = :userId', { userId }) + .orderBy('messages.created_at', 'ASC') + .getOne(); + }, + + async findChatsForConnection(connectionId: string, userId: string): Promise { + return await this.createQueryBuilder('schema_change_chat') + .where('schema_change_chat.connection_id = :connectionId', { connectionId }) + .andWhere('schema_change_chat.user_id = :userId', { userId }) + .orderBy('schema_change_chat.created_at', 'DESC') + .getMany(); + }, + + async createChatForUser(userId: string, connectionId: string, name?: string): Promise { + const newChat = this.create({ + user_id: userId, + connection_id: connectionId, + name: name || null, + }); + return await this.save(newChat); + }, + + async updateChatName(chatId: string, name: string): Promise { + await this.createQueryBuilder() + .update(SchemaChangeChatEntity) + .set({ name }) + .where('id = :chatId', { chatId }) + .execute(); + }, + + async updateLastBatchId(chatId: string, batchId: string): Promise { + await this.createQueryBuilder() + .update(SchemaChangeChatEntity) + .set({ last_batch_id: batchId }) + .where('id = :chatId', { chatId }) + .execute(); + }, +}; diff --git a/backend/src/entities/table-schema/schema-change-chat/schema-change-chat/repository/schema-change-chat-repository.interface.ts b/backend/src/entities/table-schema/schema-change-chat/schema-change-chat/repository/schema-change-chat-repository.interface.ts new file mode 100644 index 000000000..5f32e2464 --- /dev/null +++ b/backend/src/entities/table-schema/schema-change-chat/schema-change-chat/repository/schema-change-chat-repository.interface.ts @@ -0,0 +1,10 @@ +import { SchemaChangeChatEntity } from '../schema-change-chat.entity.js'; + +export interface ISchemaChangeChatRepository { + findChatByIdAndUserId(chatId: string, userId: string): Promise; + findChatWithMessagesByIdAndUserId(chatId: string, userId: string): Promise; + findChatsForConnection(connectionId: string, userId: string): Promise; + createChatForUser(userId: string, connectionId: string, name?: string): Promise; + updateChatName(chatId: string, name: string): Promise; + updateLastBatchId(chatId: string, batchId: string): Promise; +} diff --git a/backend/src/entities/table-schema/schema-change-chat/schema-change-chat/schema-change-chat.entity.ts b/backend/src/entities/table-schema/schema-change-chat/schema-change-chat/schema-change-chat.entity.ts new file mode 100644 index 000000000..01ffed51e --- /dev/null +++ b/backend/src/entities/table-schema/schema-change-chat/schema-change-chat/schema-change-chat.entity.ts @@ -0,0 +1,52 @@ +import { + Column, + CreateDateColumn, + Entity, + JoinColumn, + ManyToOne, + OneToMany, + PrimaryGeneratedColumn, + Relation, + UpdateDateColumn, +} from 'typeorm'; +import { ConnectionEntity } from '../../../connection/connection.entity.js'; +import { UserEntity } from '../../../user/user.entity.js'; +import { SchemaChangeChatMessageEntity } from '../schema-change-chat-message/schema-change-chat-message.entity.js'; + +@Entity('schema_change_chat') +export class SchemaChangeChatEntity { + @PrimaryGeneratedColumn('uuid') + id: string; + + @Column({ default: null }) + name: string; + + @CreateDateColumn({ type: 'timestamp', default: () => 'CURRENT_TIMESTAMP' }) + created_at: Date; + + @UpdateDateColumn({ type: 'timestamp', nullable: true, default: null }) + updated_at: Date; + + @ManyToOne(() => UserEntity, { onDelete: 'CASCADE' }) + @JoinColumn({ name: 'user_id' }) + user: Relation; + + @Column() + user_id: string; + + @ManyToOne(() => ConnectionEntity, { onDelete: 'CASCADE' }) + @JoinColumn({ name: 'connection_id' }) + connection: Relation; + + @Column({ type: 'varchar', length: 38 }) + connection_id: string; + + @Column({ type: 'uuid', nullable: true, default: null }) + last_batch_id: string | null; + + @OneToMany( + () => SchemaChangeChatMessageEntity, + (message) => message.chat, + ) + messages: Relation[]; +} diff --git a/backend/src/entities/table-schema/table-schema.controller.ts b/backend/src/entities/table-schema/table-schema.controller.ts index b768fb808..891f1b142 100644 --- a/backend/src/entities/table-schema/table-schema.controller.ts +++ b/backend/src/entities/table-schema/table-schema.controller.ts @@ -71,7 +71,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.', + '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.', }) @ApiParam({ name: 'connectionId', type: String }) @ApiBody({ type: GenerateSchemaChangeDto }) @@ -91,6 +91,7 @@ export class TableSchemaController { userPrompt: body.userPrompt, userId, masterPassword, + threadId: body.threadId ?? null, }); } diff --git a/backend/src/entities/table-schema/table-schema.module.ts b/backend/src/entities/table-schema/table-schema.module.ts index 21c14e894..36eef297d 100644 --- a/backend/src/entities/table-schema/table-schema.module.ts +++ b/backend/src/entities/table-schema/table-schema.module.ts @@ -8,6 +8,8 @@ import { SchemaChangeOwnershipGuard } from '../../guards/schema-change-ownership import { ConnectionEntity } from '../connection/connection.entity.js'; import { LogOutEntity } from '../log-out/log-out.entity.js'; import { UserEntity } from '../user/user.entity.js'; +import { SchemaChangeChatEntity } from './schema-change-chat/schema-change-chat/schema-change-chat.entity.js'; +import { SchemaChangeChatMessageEntity } from './schema-change-chat/schema-change-chat-message/schema-change-chat-message.entity.js'; import { TableSchemaController } from './table-schema.controller.js'; import { TableSchemaChangeEntity } from './table-schema-change.entity.js'; import { ApproveAndApplySchemaChangeUseCase } from './use-cases/approve-and-apply-schema-change.use-case.js'; @@ -22,7 +24,16 @@ import { RollbackBatchSchemaChangesUseCase } from './use-cases/rollback-batch-sc import { RollbackSchemaChangeUseCase } from './use-cases/rollback-schema-change.use-case.js'; @Module({ - imports: [TypeOrmModule.forFeature([TableSchemaChangeEntity, ConnectionEntity, UserEntity, LogOutEntity])], + imports: [ + TypeOrmModule.forFeature([ + TableSchemaChangeEntity, + SchemaChangeChatEntity, + SchemaChangeChatMessageEntity, + ConnectionEntity, + UserEntity, + LogOutEntity, + ]), + ], providers: [ { provide: BaseType.GLOBAL_DB_CONTEXT, diff --git a/backend/src/entities/table-schema/use-cases/generate-schema-change.use-case.ts b/backend/src/entities/table-schema/use-cases/generate-schema-change.use-case.ts index b8959cf1d..adf8c0a71 100644 --- a/backend/src/entities/table-schema/use-cases/generate-schema-change.use-case.ts +++ b/backend/src/entities/table-schema/use-cases/generate-schema-change.use-case.ts @@ -1,6 +1,8 @@ +import { BaseMessage } from '@langchain/core/messages'; import { BadRequestException, Inject, Injectable, Logger, NotFoundException, Scope } 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 Sentry from '@sentry/minimal'; import crypto from 'crypto'; import { AIProviderType } from '../../../ai-core/interfaces/ai-service.interface.js'; import { AICoreService } from '../../../ai-core/services/ai-core.service.js'; @@ -9,6 +11,7 @@ import AbstractUseCase from '../../../common/abstract-use.case.js'; import { IGlobalDatabaseContext } from '../../../common/application/global-database-context.interface.js'; import { BaseType } from '../../../common/data-injection.tokens.js'; import { Messages } from '../../../exceptions/text/messages.js'; +import { MessageRole } from '../../ai/ai-conversation-history/ai-chat-messages/message-role.enum.js'; import { runSchemaChangeAiLoop } from '../ai/run-schema-change-ai-loop.js'; import { buildSchemaChangePrompt } from '../ai/schema-change-prompts.js'; import { @@ -20,6 +23,7 @@ import { } from '../ai/schema-change-tools.js'; import { GenerateSchemaChangeDs } from '../application/data-structures/generate-schema-change.ds.js'; import { SchemaChangeBatchResponseDto } from '../application/data-transfer-objects/schema-change-batch-response.dto.js'; +import { SchemaChangeChatEntity } from '../schema-change-chat/schema-change-chat/schema-change-chat.entity.js'; import { TableSchemaChangeEntity } from '../table-schema-change.entity.js'; import { isDynamoDbSchemaChangeType, @@ -58,7 +62,7 @@ export class GenerateSchemaChangeUseCase } protected async implementation(inputData: GenerateSchemaChangeDs): Promise { - const { connectionId, userPrompt, userId, masterPassword } = inputData; + const { connectionId, userPrompt, userId, masterPassword, threadId } = inputData; const connection = await this._dbContext.connectionRepository.findAndDecryptConnection( connectionId, @@ -77,12 +81,14 @@ export class GenerateSchemaChangeUseCase throw new BadRequestException(Messages.AI_REQUESTS_NOT_ALLOWED); } + const { chat, isNewChat } = await this.resolveChat(threadId ?? null, userId, connectionId); + const dao = getDataAccessObject(connection); const tableList = await dao.getTablesFromDB(); const tableNames = tableList.map((t) => t.tableName); const systemPrompt = buildSchemaChangePrompt(connectionType, tableNames, connection.schema ?? null); - const messages = new MessageBuilder().system(systemPrompt).human(userPrompt).build(); + const messages = await this.buildMessagesWithHistory(systemPrompt, userPrompt, chat.id, isNewChat); const tools = isMongoDialect(connectionType) ? createMongoSchemaChangeTools() : isDynamoDbDialect(connectionType) @@ -91,6 +97,14 @@ export class GenerateSchemaChangeUseCase ? createElasticsearchSchemaChangeTools() : createSchemaChangeTools(); + await this._dbContext.schemaChangeChatMessageRepository.saveMessage(chat.id, userPrompt, MessageRole.user); + + if (isNewChat) { + this.generateAndUpdateChatName(chat.id, userPrompt).catch((error) => { + Sentry.captureException(error); + }); + } + let proposals: ProposeSchemaChangeArgs[]; try { const result = await runSchemaChangeAiLoop({ @@ -141,12 +155,101 @@ export class GenerateSchemaChangeUseCase const saved = await this._dbContext.tableSchemaChangeRepository.createPendingBatch(items); saved.sort((a, b) => a.orderInBatch - b.orderInBatch); + await this._dbContext.schemaChangeChatMessageRepository.saveMessage( + chat.id, + this.serializeAssistantTurn(proposals), + MessageRole.ai, + batchId, + ); + await this._dbContext.schemaChangeChatRepository.updateLastBatchId(chat.id, batchId); + return { batchId, + threadId: chat.id, changes: saved.map(mapSchemaChangeToResponseDto), }; } + private async resolveChat( + threadId: string | null, + userId: string, + connectionId: string, + ): Promise<{ chat: SchemaChangeChatEntity; isNewChat: boolean }> { + if (threadId) { + const existing = await this._dbContext.schemaChangeChatRepository.findChatByIdAndUserId(threadId, userId); + if (existing) { + if (existing.connection_id !== connectionId) { + throw new BadRequestException('Provided threadId belongs to a different connection.'); + } + return { chat: existing, isNewChat: false }; + } + } + const created = await this._dbContext.schemaChangeChatRepository.createChatForUser(userId, connectionId); + return { chat: created, isNewChat: true }; + } + + private async buildMessagesWithHistory( + systemPrompt: string, + userMessage: string, + chatId: string, + isNewChat: boolean, + ): Promise { + if (isNewChat) { + return new MessageBuilder().system(systemPrompt).human(userMessage).build(); + } + + const previousMessages = await this._dbContext.schemaChangeChatMessageRepository.findMessagesForChat(chatId); + const builder = new MessageBuilder().system(systemPrompt); + + for (const msg of previousMessages) { + if (msg.role === MessageRole.user) { + builder.human(msg.message); + } else if (msg.role === MessageRole.ai) { + builder.ai(msg.message); + } + } + + builder.human(userMessage); + return builder.build(); + } + + private serializeAssistantTurn(proposals: ProposeSchemaChangeArgs[]): string { + return proposals + .map((p, i) => { + const summary = p.summary?.trim() || '(no summary)'; + return `${i + 1}. [${p.changeType}] ${p.targetTableName} — ${summary}`; + }) + .join('\n'); + } + + private async generateAndUpdateChatName(chatId: string, userMessage: string): Promise { + try { + const CHAT_NAME_GENERATION_PROMPT = `Generate a very short, concise title (max 5-6 words) for a database schema-change conversation based on the user's first request. +The title should capture the main intent (e.g. "Add products table", "Rename users column"). +Respond ONLY with the title, no quotes, no explanation. +User request: `; + const prompt = CHAT_NAME_GENERATION_PROMPT + userMessage; + const messages = new MessageBuilder().human(prompt).build(); + + let generatedName = ''; + const stream = await this.aiCoreService.streamChatWithToolsAndProvider(this.provider, messages, []); + + for await (const chunk of stream) { + if (chunk.type === 'text' && chunk.content) { + generatedName += chunk.content; + } + } + + generatedName = generatedName.trim().slice(0, 100); + + if (generatedName) { + await this._dbContext.schemaChangeChatRepository.updateChatName(chatId, generatedName); + } + } catch (error) { + Sentry.captureException(error); + } + } + private validateProposal( proposal: ProposeSchemaChangeArgs, connectionType: ConnectionTypesEnum, diff --git a/backend/src/migrations/1778767036234-AddSchemaChangeChatEntities.ts b/backend/src/migrations/1778767036234-AddSchemaChangeChatEntities.ts new file mode 100644 index 000000000..ba00cb234 --- /dev/null +++ b/backend/src/migrations/1778767036234-AddSchemaChangeChatEntities.ts @@ -0,0 +1,39 @@ +import { MigrationInterface, QueryRunner } from 'typeorm'; + +export class AddSchemaChangeChatEntities1778767036234 implements MigrationInterface { + name = 'AddSchemaChangeChatEntities1778767036234'; + + public async up(queryRunner: QueryRunner): Promise { + await queryRunner.query( + `CREATE TABLE "schema_change_chat" ("id" uuid NOT NULL DEFAULT uuid_generate_v4(), "name" character varying, "created_at" TIMESTAMP NOT NULL DEFAULT now(), "updated_at" TIMESTAMP DEFAULT now(), "user_id" uuid NOT NULL, "connection_id" character varying(38) NOT NULL, "last_batch_id" uuid, CONSTRAINT "PK_60082e3e240c265fc043290381d" PRIMARY KEY ("id"))`, + ); + await queryRunner.query( + `CREATE TYPE "public"."schema_change_chat_message_role_enum" AS ENUM('user', 'ai', 'system')`, + ); + await queryRunner.query( + `CREATE TABLE "schema_change_chat_message" ("id" uuid NOT NULL DEFAULT uuid_generate_v4(), "message" text, "role" "public"."schema_change_chat_message_role_enum", "batch_id" uuid, "created_at" TIMESTAMP NOT NULL DEFAULT now(), "updated_at" TIMESTAMP DEFAULT now(), "chat_id" uuid NOT NULL, CONSTRAINT "PK_5984cdb248fa9c2f55f5a19022c" PRIMARY KEY ("id"))`, + ); + await queryRunner.query(`ALTER TABLE "ai_chat_message" DROP COLUMN "response_id"`); + await queryRunner.query( + `ALTER TABLE "schema_change_chat" ADD CONSTRAINT "FK_4dbf7dad457505747189fb98d7e" FOREIGN KEY ("user_id") REFERENCES "user"("id") ON DELETE CASCADE ON UPDATE NO ACTION`, + ); + await queryRunner.query( + `ALTER TABLE "schema_change_chat" ADD CONSTRAINT "FK_9f9acf0578fcf239576640d7b7b" FOREIGN KEY ("connection_id") REFERENCES "connection"("id") ON DELETE CASCADE ON UPDATE NO ACTION`, + ); + await queryRunner.query( + `ALTER TABLE "schema_change_chat_message" ADD CONSTRAINT "FK_32825f4780664738f60fa75cd50" FOREIGN KEY ("chat_id") REFERENCES "schema_change_chat"("id") ON DELETE CASCADE ON UPDATE NO ACTION`, + ); + } + + public async down(queryRunner: QueryRunner): Promise { + await queryRunner.query( + `ALTER TABLE "schema_change_chat_message" DROP CONSTRAINT "FK_32825f4780664738f60fa75cd50"`, + ); + await queryRunner.query(`ALTER TABLE "schema_change_chat" DROP CONSTRAINT "FK_9f9acf0578fcf239576640d7b7b"`); + await queryRunner.query(`ALTER TABLE "schema_change_chat" DROP CONSTRAINT "FK_4dbf7dad457505747189fb98d7e"`); + await queryRunner.query(`ALTER TABLE "ai_chat_message" ADD "response_id" character varying(255)`); + await queryRunner.query(`DROP TABLE "schema_change_chat_message"`); + await queryRunner.query(`DROP TYPE "public"."schema_change_chat_message_role_enum"`); + await queryRunner.query(`DROP TABLE "schema_change_chat"`); + } +} diff --git a/backend/test/ava-tests/non-saas-tests/non-saas-table-schema-postgres-e2e.test.ts b/backend/test/ava-tests/non-saas-tests/non-saas-table-schema-postgres-e2e.test.ts index efacfa08a..88fc9856d 100644 --- a/backend/test/ava-tests/non-saas-tests/non-saas-table-schema-postgres-e2e.test.ts +++ b/backend/test/ava-tests/non-saas-tests/non-saas-table-schema-postgres-e2e.test.ts @@ -6,9 +6,14 @@ import test from 'ava'; import { ValidationError } from 'class-validator'; import cookieParser from 'cookie-parser'; import request from 'supertest'; +import { DataSource } from 'typeorm'; import { AICoreService } from '../../../src/ai-core/services/ai-core.service.js'; import { ApplicationModule } from '../../../src/app.module.js'; +import { BaseType } from '../../../src/common/data-injection.tokens.js'; +import { MessageRole } from '../../../src/entities/ai/ai-conversation-history/ai-chat-messages/message-role.enum.js'; import { WinstonLogger } from '../../../src/entities/logging/winston-logger.js'; +import { SchemaChangeChatEntity } from '../../../src/entities/table-schema/schema-change-chat/schema-change-chat/schema-change-chat.entity.js'; +import { SchemaChangeChatMessageEntity } from '../../../src/entities/table-schema/schema-change-chat/schema-change-chat-message/schema-change-chat-message.entity.js'; import { SchemaChangeStatusEnum, SchemaChangeTypeEnum, @@ -981,3 +986,153 @@ test.serial('multi-table batch: rollback removes tables in reverse and creates l const auditRows = list.data.filter((c: any) => c.changeType === SchemaChangeTypeEnum.ROLLBACK); t.is(auditRows.length, 2); }); + +test.serial('chat: generate without threadId creates a chat thread and persists user + ai turns', async (t) => { + const { token } = await registerUserAndReturnUserInfo(app); + const connectionId = await createConnection(token); + const tableName = `ra_t_${faker.string.alphanumeric(8).toLowerCase()}`; + testTables.push(tableName); + + nextProposal = { + forwardSql: `CREATE TABLE "${tableName}" (id SERIAL PRIMARY KEY)`, + rollbackSql: `DROP TABLE "${tableName}"`, + changeType: SchemaChangeTypeEnum.CREATE_TABLE, + targetTableName: tableName, + isReversible: true, + summary: `Create ${tableName}`, + reasoning: 'Initial turn.', + }; + + const userPrompt = `create a ${tableName} table`; + const generateResp = await request(app.getHttpServer()) + .post(`/table-schema/${connectionId}/generate`) + .set('Cookie', token) + .set('Content-Type', 'application/json') + .send({ userPrompt }); + t.is(generateResp.status, 201); + const batch = JSON.parse(generateResp.text); + t.truthy(batch.threadId); + t.truthy(batch.batchId); + + const dataSource = app.get(BaseType.DATA_SOURCE); + const chatRepo = dataSource.getRepository(SchemaChangeChatEntity); + const messageRepo = dataSource.getRepository(SchemaChangeChatMessageEntity); + + const chat = await chatRepo.findOne({ where: { id: batch.threadId } }); + t.truthy(chat); + t.is(chat!.connection_id, connectionId); + t.is(chat!.last_batch_id, batch.batchId); + + const messages = await messageRepo.find({ where: { chat_id: batch.threadId }, order: { created_at: 'ASC' } }); + t.is(messages.length, 2); + t.is(messages[0].role, MessageRole.user); + t.is(messages[0].message, userPrompt); + t.is(messages[1].role, MessageRole.ai); + t.is(messages[1].batch_id, batch.batchId); + t.true(messages[1].message.includes(tableName)); +}); + +test.serial('chat: passing threadId continues the conversation and appends turns', async (t) => { + const { token } = await registerUserAndReturnUserInfo(app); + const connectionId = await createConnection(token); + const tableName = `ra_t_${faker.string.alphanumeric(8).toLowerCase()}`; + testTables.push(tableName); + + nextProposal = { + forwardSql: `CREATE TABLE "${tableName}" (id SERIAL PRIMARY KEY)`, + rollbackSql: `DROP TABLE "${tableName}"`, + changeType: SchemaChangeTypeEnum.CREATE_TABLE, + targetTableName: tableName, + isReversible: true, + summary: `Create ${tableName}`, + reasoning: 'First turn.', + }; + const firstResp = await request(app.getHttpServer()) + .post(`/table-schema/${connectionId}/generate`) + .set('Cookie', token) + .set('Content-Type', 'application/json') + .send({ userPrompt: `create a ${tableName} table` }); + t.is(firstResp.status, 201); + const firstBatch = JSON.parse(firstResp.text); + const threadId = firstBatch.threadId; + t.truthy(threadId); + + const indexName = `${tableName}_idx`; + nextProposal = { + forwardSql: `CREATE INDEX "${indexName}" ON "${tableName}" (id)`, + rollbackSql: `DROP INDEX "${indexName}"`, + changeType: SchemaChangeTypeEnum.ADD_INDEX, + targetTableName: tableName, + isReversible: true, + summary: `Add index ${indexName}`, + reasoning: 'Follow-up.', + }; + const secondResp = await request(app.getHttpServer()) + .post(`/table-schema/${connectionId}/generate`) + .set('Cookie', token) + .set('Content-Type', 'application/json') + .send({ userPrompt: 'now add an index on id', threadId }); + t.is(secondResp.status, 201); + const secondBatch = JSON.parse(secondResp.text); + t.is(secondBatch.threadId, threadId); + t.not(secondBatch.batchId, firstBatch.batchId); + + const dataSource = app.get(BaseType.DATA_SOURCE); + const chatRepo = dataSource.getRepository(SchemaChangeChatEntity); + const messageRepo = dataSource.getRepository(SchemaChangeChatMessageEntity); + + const chat = await chatRepo.findOne({ where: { id: threadId } }); + t.is(chat!.last_batch_id, secondBatch.batchId); + + const messages = await messageRepo.find({ where: { chat_id: threadId }, order: { created_at: 'ASC' } }); + t.is(messages.length, 4); + t.is(messages[0].role, MessageRole.user); + t.is(messages[1].role, MessageRole.ai); + t.is(messages[1].batch_id, firstBatch.batchId); + t.is(messages[2].role, MessageRole.user); + t.is(messages[2].message, 'now add an index on id'); + t.is(messages[3].role, MessageRole.ai); + t.is(messages[3].batch_id, secondBatch.batchId); +}); + +test.serial('chat: threadId from a different connection is rejected with 400', async (t) => { + const { token } = await registerUserAndReturnUserInfo(app); + const connectionIdA = await createConnection(token); + const connectionIdB = await createConnection(token); + const tableName = `ra_t_${faker.string.alphanumeric(8).toLowerCase()}`; + testTables.push(tableName); + + nextProposal = { + forwardSql: `CREATE TABLE "${tableName}" (id SERIAL PRIMARY KEY)`, + rollbackSql: `DROP TABLE "${tableName}"`, + changeType: SchemaChangeTypeEnum.CREATE_TABLE, + targetTableName: tableName, + isReversible: true, + summary: `Create ${tableName}`, + reasoning: 'For cross-connection guard.', + }; + const firstResp = await request(app.getHttpServer()) + .post(`/table-schema/${connectionIdA}/generate`) + .set('Cookie', token) + .set('Content-Type', 'application/json') + .send({ userPrompt: `create a ${tableName} table` }); + t.is(firstResp.status, 201); + const threadId = JSON.parse(firstResp.text).threadId; + t.truthy(threadId); + + nextProposal = { + forwardSql: `CREATE TABLE "${tableName}_b" (id SERIAL PRIMARY KEY)`, + rollbackSql: `DROP TABLE "${tableName}_b"`, + changeType: SchemaChangeTypeEnum.CREATE_TABLE, + targetTableName: `${tableName}_b`, + isReversible: true, + summary: 'Should not be reached', + reasoning: 'Should not be reached', + }; + const crossResp = await request(app.getHttpServer()) + .post(`/table-schema/${connectionIdB}/generate`) + .set('Cookie', token) + .set('Content-Type', 'application/json') + .send({ userPrompt: 'create another table', threadId }); + t.is(crossResp.status, 400); +});