| 1 | import type { BeforeApplicationShutdown } from '@nestjs/common' |
| 2 | import type { DraftGenerationData } from '@yikart/aitoearn-queue' |
| 3 | import type { Job, Worker } from 'bullmq' |
| 4 | import { WorkerHost } from '@nestjs/bullmq' |
| 5 | import { Logger } from '@nestjs/common' |
| 6 | import { QueueName, QueueProcessor } from '@yikart/aitoearn-queue' |
| 7 | import { getErrorMessage } from '@yikart/common' |
| 8 | import { AiLogRepository, AiLogStatus } from '@yikart/mongodb' |
| 9 | import { UnrecoverableError } from 'bullmq' |
| 10 | import { config } from '../../config' |
| 11 | import { DraftGenerationError, DraftGenerationService } from './draft-generation.service' |
| 12 | |
| 13 | interface ActiveDraftGenerationJob { |
| 14 | job: Job<DraftGenerationData> |
| 15 | token?: string |
| 16 | } |
| 17 | |
| 18 | abstract class DraftGenerationConsumerBase extends WorkerHost implements BeforeApplicationShutdown { |
| 19 | private readonly logger: Logger |
| 20 | private readonly activeJobs = new Map<string, ActiveDraftGenerationJob>() |
| 21 | |
| 22 | constructor( |
| 23 | private readonly draftGenerationService: DraftGenerationService, |
| 24 | private readonly aiLogRepository: AiLogRepository, |
| 25 | loggerName: string, |
| 26 | ) { |
| 27 | super() |
| 28 | this.logger = new Logger(loggerName) |
| 29 | } |
| 30 | |
| 31 | override async process(job: Job<DraftGenerationData>, token?: string): Promise<void> { |
| 32 | const { aiLogId, userId, userType, groupId, version } = job.data |
| 33 | this.activeJobs.set(this.getActiveJobKey(job), { job, token }) |
| 34 | |
| 35 | try { |
| 36 | if (job.attemptsMade > 0) { |
| 37 | const aiLog = await this.aiLogRepository.getById(aiLogId) |
| 38 | if (aiLog?.status === AiLogStatus.Success) { |
| 39 | this.logger.log({ aiLogId, attemptsMade: job.attemptsMade }, 'Skipping retry: AiLog already succeeded') |
| 40 | return |
| 41 | } |
| 42 | this.logger.log({ aiLogId, attemptsMade: job.attemptsMade }, 'Retrying draft generation') |
| 43 | await this.aiLogRepository.updateById(aiLogId, { |
| 44 | $set: { status: AiLogStatus.Generating }, |
| 45 | $unset: { errorMessage: '' }, |
| 46 | }) |
| 47 | } |
| 48 | else { |
| 49 | const aiLog = await this.aiLogRepository.getById(aiLogId) |
| 50 | if (aiLog && aiLog.status !== AiLogStatus.Generating) { |
| 51 | this.logger.warn({ aiLogId, status: aiLog.status }, 'Skipping draft generation because AiLog is no longer generating') |
| 52 | return |
| 53 | } |
| 54 | } |
| 55 | |
| 56 | if (version === 'v2-image-text') { |
| 57 | const { prompt, captionPrompt, imageUrls, imageModel, imageCount, imageSize, aspectRatio, imageTextDraftType, platforms, plannerModel, disableMemory } = job.data |
| 58 | this.logger.log( |
| 59 | { aiLogId, imageModel, imageCount, aspectRatio, imageUrlsCount: imageUrls?.length ?? 0, promptLength: prompt?.length ?? 0, draftType: imageTextDraftType }, |
| 60 | 'Processing v2-image-text generation', |
| 61 | ) |
| 62 | await this.draftGenerationService.generateContentImageText(aiLogId, userId, userType, groupId, { |
| 63 | prompt: prompt ?? '', |
| 64 | captionPrompt, |
| 65 | imageUrls, |
| 66 | imageModel: imageModel ?? 'gemini-3.1-flash-image-preview', |
| 67 | imageCount: imageCount ?? 3, |
| 68 | imageSize, |
| 69 | aspectRatio, |
| 70 | draftType: imageTextDraftType, |
| 71 | platforms, |
| 72 | plannerModel, |
| 73 | disableMemory, |
| 74 | }) |
| 75 | this.logger.log({ aiLogId }, 'v2-image-text generation completed') |
| 76 | } |
| 77 | else if (version === 'v2') { |
| 78 | const { prompt, captionPrompt, imageUrls, model, duration, resolution, aspectRatio, videoUrls, audioUrls, draftType, platforms, plannerModel, disableMemory } = job.data |
| 79 | await this.draftGenerationService.generateContentV2(aiLogId, userId, userType, groupId, { |
| 80 | prompt, |
| 81 | captionPrompt, |
| 82 | imageUrls, |
| 83 | model, |
| 84 | duration, |
| 85 | resolution, |
| 86 | aspectRatio, |
| 87 | videoUrls, |
| 88 | audioUrls, |
| 89 | draftType, |
| 90 | platforms, |
| 91 | plannerModel, |
| 92 | disableMemory, |
| 93 | }) |
| 94 | this.logger.log({ aiLogId }, 'v2 generation completed') |
| 95 | } |
| 96 | else { |
| 97 | const unsupportedVersion = String((job.data as { version?: string }).version ?? 'missing') |
| 98 | throw new UnrecoverableError(`Unsupported draft generation version: ${unsupportedVersion}`) |
| 99 | } |
| 100 | } |
| 101 | catch (error) { |
| 102 | const originalError = error instanceof DraftGenerationError ? (error.cause ?? error) : error |
| 103 | const errorMessage = getErrorMessage(originalError) |
| 104 | const versionLabel = String((job.data as { version?: string }).version ?? 'missing') |
| 105 | |
| 106 | this.logger.error( |
| 107 | originalError, |
| 108 | `DraftGeneration failed (version=${versionLabel}, aiLogId=${aiLogId}, userId=${userId})`, |
| 109 | ) |
| 110 | |
| 111 | await this.aiLogRepository.updateById(aiLogId, { |
| 112 | $set: { |
| 113 | status: AiLogStatus.Failed, |
| 114 | errorMessage, |
| 115 | }, |
| 116 | }) |
| 117 | |
| 118 | throw error |
| 119 | } |
| 120 | finally { |
| 121 | this.activeJobs.delete(this.getActiveJobKey(job)) |
| 122 | } |
| 123 | } |
| 124 | |
| 125 | async beforeApplicationShutdown() { |
| 126 | const worker = this.getInitializedWorker() |
| 127 | if (!worker) { |
| 128 | return |
| 129 | } |
| 130 | |
| 131 | await worker.pause(true) |
| 132 | |
| 133 | const activeJobs = [...this.activeJobs.values()] |
| 134 | if (activeJobs.length === 0) { |
| 135 | await worker.close(true) |
| 136 | return |
| 137 | } |
| 138 | |
| 139 | this.logger.warn({ count: activeJobs.length }, 'Failing active draft generation jobs before shutdown') |
| 140 | |
| 141 | for (const activeJob of activeJobs) { |
| 142 | await this.failActiveJob(activeJob) |
| 143 | } |
| 144 | |
| 145 | await worker.close(true) |
| 146 | } |
| 147 | |
| 148 | private async failActiveJob(activeJob: ActiveDraftGenerationJob) { |
| 149 | const { job, token } = activeJob |
| 150 | const { aiLogId } = job.data |
| 151 | const errorMessage = 'Application is shutting down; draft generation will retry' |
| 152 | |
| 153 | if (!token) { |
| 154 | this.logger.warn({ aiLogId, jobId: job.id }, 'Cannot fail active draft generation job without lock token') |
| 155 | return |
| 156 | } |
| 157 | |
| 158 | try { |
| 159 | await job.moveToFailed(new Error(errorMessage), token, false) |
| 160 | await this.aiLogRepository.updateById(aiLogId, { |
| 161 | $set: { |
| 162 | status: AiLogStatus.Failed, |
| 163 | errorMessage, |
| 164 | }, |
| 165 | }) |
| 166 | } |
| 167 | catch (error) { |
| 168 | this.logger.error({ error, aiLogId, jobId: job.id }, 'Failed to move active draft generation job to failed state') |
| 169 | } |
| 170 | } |
| 171 | |
| 172 | private getInitializedWorker(): Worker | undefined { |
| 173 | try { |
| 174 | return this.worker |
| 175 | } |
| 176 | catch (error) { |
| 177 | this.logger.warn({ error }, 'Draft generation worker is not initialized') |
| 178 | return undefined |
| 179 | } |
| 180 | } |
| 181 | |
| 182 | private getActiveJobKey(job: Job<DraftGenerationData>) { |
| 183 | return job.id ?? job.data.aiLogId |
| 184 | } |
| 185 | } |
| 186 | |
| 187 | @QueueProcessor(QueueName.DraftGeneration, { concurrency: 60 }) |
| 188 | export class DraftGenerationConsumer extends DraftGenerationConsumerBase { |
| 189 | constructor( |
| 190 | draftGenerationService: DraftGenerationService, |
| 191 | aiLogRepository: AiLogRepository, |
| 192 | ) { |
| 193 | super(draftGenerationService, aiLogRepository, DraftGenerationConsumer.name) |
| 194 | } |
| 195 | |
| 196 | override async process(job: Job<DraftGenerationData>, token?: string): Promise<void> { |
| 197 | return super.process(job, token) |
| 198 | } |
| 199 | } |
| 200 | |
| 201 | @QueueProcessor(QueueName.DraftGenerationLowPriority, { |
| 202 | concurrency: config.ai.draftGeneration.queue.lowPriorityConcurrency, |
| 203 | }) |
| 204 | export class DraftGenerationLowPriorityConsumer extends DraftGenerationConsumerBase { |
| 205 | constructor( |
| 206 | draftGenerationService: DraftGenerationService, |
| 207 | aiLogRepository: AiLogRepository, |
| 208 | ) { |
| 209 | super(draftGenerationService, aiLogRepository, DraftGenerationLowPriorityConsumer.name) |
| 210 | } |
| 211 | |
| 212 | override async process(job: Job<DraftGenerationData>, token?: string): Promise<void> { |
| 213 | return super.process(job, token) |
| 214 | } |
| 215 | } |
| 216 |