import { randomUUID } from 'crypto'; import { eq, desc, count, and, isNull, inArray } from 'drizzle-orm'; import { db } from '@/db'; import { generations, flows, images } from '@banatie/database'; import type { Generation, NewGeneration, GenerationWithRelations, GenerationFilters, } from '@/types/models'; import { ImageService } from './ImageService'; import { AliasService } from './AliasService'; import { ImageGenService } from '../ImageGenService'; import { StorageFactory } from '../StorageFactory'; import { buildWhereClause, buildEqCondition } from '@/utils/helpers'; import { ERROR_MESSAGES, GENERATION_LIMITS } from '@/utils/constants'; import { extractAliasesFromPrompt } from '@/utils/validators'; import type { ReferenceImage } from '@/types/api'; export interface CreateGenerationParams { projectId: string; apiKeyId: string; prompt: string; referenceImages?: string[] | undefined; // Aliases to resolve aspectRatio?: string | undefined; flowId?: string | undefined; alias?: string | undefined; flowAlias?: string | undefined; autoEnhance?: boolean | undefined; enhancedPrompt?: string | undefined; meta?: Record | undefined; requestId?: string | undefined; } export class GenerationService { private imageService: ImageService; private aliasService: AliasService; private imageGenService: ImageGenService; constructor() { this.imageService = new ImageService(); this.aliasService = new AliasService(); const geminiApiKey = process.env['GEMINI_API_KEY']; if (!geminiApiKey) { throw new Error('GEMINI_API_KEY environment variable is required'); } this.imageGenService = new ImageGenService(geminiApiKey); } async create(params: CreateGenerationParams): Promise { const startTime = Date.now(); // Auto-detect aliases from prompt and merge with manual references const autoDetectedAliases = extractAliasesFromPrompt(params.prompt); const manualReferences = params.referenceImages || []; // Merge: manual references first, then auto-detected (remove duplicates) const allReferences = Array.from(new Set([...manualReferences, ...autoDetectedAliases])); // FlowId logic (Section 10.1 - UPDATED FOR LAZY PATTERN): // - If undefined → generate UUID for pendingFlowId, flowId = null (lazy) // - If null → flowId = null, pendingFlowId = null (explicitly no flow) // - If string → flowId = string, pendingFlowId = null (use provided, create if needed) let finalFlowId: string | null; let pendingFlowId: string | null = null; if (params.flowId === undefined) { // Lazy pattern: defer flow creation until needed pendingFlowId = randomUUID(); finalFlowId = null; } else if (params.flowId === null) { // Explicitly no flow finalFlowId = null; pendingFlowId = null; } else { // Specific flowId provided - ensure flow exists (eager creation) finalFlowId = params.flowId; pendingFlowId = null; // Check if flow exists, create if not const existingFlow = await db.query.flows.findFirst({ where: eq(flows.id, finalFlowId), }); if (!existingFlow) { await db.insert(flows).values({ id: finalFlowId, projectId: params.projectId, aliases: {}, meta: {}, }); // Link any pending generations to this new flow await this.linkPendingGenerationsToFlow(finalFlowId, params.projectId); } } // Prompt semantics (Section 2.1): // - originalPrompt: ALWAYS contains user's original input // - prompt: Enhanced version if autoEnhance=true, otherwise same as originalPrompt const usedPrompt = params.enhancedPrompt || params.prompt; const preservedOriginal = params.prompt; // Always store original const generationRecord: NewGeneration = { projectId: params.projectId, flowId: finalFlowId, pendingFlowId: pendingFlowId, apiKeyId: params.apiKeyId, status: 'pending', prompt: usedPrompt, // Prompt actually used for generation originalPrompt: preservedOriginal, // User's original (only if enhanced) aspectRatio: params.aspectRatio || GENERATION_LIMITS.DEFAULT_ASPECT_RATIO, referencedImages: null, requestId: params.requestId || null, meta: params.meta || {}, }; const [generation] = await db .insert(generations) .values(generationRecord) .returning(); if (!generation) { throw new Error('Failed to create generation record'); } try { await this.updateStatus(generation.id, 'processing'); let referenceImageBuffers: ReferenceImage[] = []; let referencedImagesMetadata: Array<{ imageId: string; alias: string }> = []; if (allReferences.length > 0) { const resolved = await this.resolveReferenceImages( allReferences, params.projectId, params.flowId ); referenceImageBuffers = resolved.buffers; referencedImagesMetadata = resolved.metadata; await db .update(generations) .set({ referencedImages: referencedImagesMetadata }) .where(eq(generations.id, generation.id)); } const genResult = await this.imageGenService.generateImage({ prompt: usedPrompt, // Use the prompt that was stored (enhanced or original) filename: `gen_${generation.id}`, referenceImages: referenceImageBuffers, aspectRatio: params.aspectRatio || GENERATION_LIMITS.DEFAULT_ASPECT_RATIO, orgId: 'default', projectId: params.projectId, meta: params.meta || {}, }); if (!genResult.success) { const processingTime = Date.now() - startTime; await this.updateStatus(generation.id, 'failed', { errorMessage: genResult.error || 'Generation failed', processingTimeMs: processingTime, }); throw new Error(genResult.error || 'Generation failed'); } const storageKey = genResult.filepath!; // TODO: Add file hash computation when we have a helper to download by storageKey const fileHash = null; const imageRecord = await this.imageService.create({ projectId: params.projectId, flowId: finalFlowId, generationId: generation.id, apiKeyId: params.apiKeyId, storageKey, storageUrl: genResult.url!, mimeType: 'image/jpeg', fileSize: genResult.size || 0, fileHash, source: 'generated', alias: null, meta: params.meta || {}, width: genResult.generatedImageData?.width ?? null, height: genResult.generatedImageData?.height ?? null, }); // Reassign project alias if provided (override behavior per Section 5.2) if (params.alias) { await this.imageService.reassignProjectAlias( params.alias, imageRecord.id, params.projectId ); } // Eager flow creation if flowAlias is provided (Section 4.2) if (params.flowAlias) { // If we have pendingFlowId, create flow and link pending generations const flowIdToUse = pendingFlowId || finalFlowId; if (!flowIdToUse) { throw new Error('Cannot create flow: no flowId available'); } // Check if flow exists, create if not const existingFlow = await db.query.flows.findFirst({ where: eq(flows.id, flowIdToUse), }); if (!existingFlow) { await db.insert(flows).values({ id: flowIdToUse, projectId: params.projectId, aliases: {}, meta: {}, }); // Link any pending generations to this new flow await this.linkPendingGenerationsToFlow(flowIdToUse, params.projectId); } await this.assignFlowAlias(flowIdToUse, params.flowAlias, imageRecord.id); } // Update flow timestamp if flow was created (either from finalFlowId or pendingFlowId converted to flow) const actualFlowId = finalFlowId || (pendingFlowId && params.flowAlias ? pendingFlowId : null); if (actualFlowId) { await db .update(flows) .set({ updatedAt: new Date() }) .where(eq(flows.id, actualFlowId)); } const processingTime = Date.now() - startTime; await this.updateStatus(generation.id, 'success', { outputImageId: imageRecord.id, processingTimeMs: processingTime, }); return await this.getByIdWithRelations(generation.id); } catch (error) { const processingTime = Date.now() - startTime; await this.updateStatus(generation.id, 'failed', { errorMessage: error instanceof Error ? error.message : 'Unknown error', processingTimeMs: processingTime, }); throw error; } } private async resolveReferenceImages( aliases: string[], projectId: string, flowId?: string ): Promise<{ buffers: ReferenceImage[]; metadata: Array<{ imageId: string; alias: string }>; }> { const resolutions = await this.aliasService.resolveMultiple(aliases, projectId, flowId); const buffers: ReferenceImage[] = []; const metadata: Array<{ imageId: string; alias: string }> = []; const storageService = await StorageFactory.getInstance(); for (const [alias, resolution] of resolutions) { if (!resolution.image) { throw new Error(`${ERROR_MESSAGES.ALIAS_NOT_FOUND}: ${alias}`); } const parts = resolution.image.storageKey.split('/'); if (parts.length < 4) { throw new Error(`Invalid storage key format: ${resolution.image.storageKey}`); } const orgId = parts[0]!; const projId = parts[1]!; const category = parts[2]! as 'uploads' | 'generated' | 'references'; const filename = parts.slice(3).join('/'); const buffer = await storageService.downloadFile( orgId, projId, category, filename ); buffers.push({ buffer, mimetype: resolution.image.mimeType, originalname: filename, }); metadata.push({ imageId: resolution.imageId, alias, }); } return { buffers, metadata }; } private async assignFlowAlias( flowId: string, flowAlias: string, imageId: string ): Promise { const flow = await db.query.flows.findFirst({ where: eq(flows.id, flowId), }); if (!flow) { throw new Error(ERROR_MESSAGES.FLOW_NOT_FOUND); } const currentAliases = (flow.aliases as Record) || {}; const updatedAliases = { ...currentAliases }; // Assign the flow alias to the image updatedAliases[flowAlias] = imageId; await db .update(flows) .set({ aliases: updatedAliases, updatedAt: new Date() }) .where(eq(flows.id, flowId)); } private async linkPendingGenerationsToFlow( flowId: string, projectId: string ): Promise { // Find all generations with pendingFlowId matching this flowId const pendingGens = await db.query.generations.findMany({ where: and( eq(generations.pendingFlowId, flowId), eq(generations.projectId, projectId) ), }); if (pendingGens.length === 0) { return; } // Update generations: set flowId and clear pendingFlowId await db .update(generations) .set({ flowId: flowId, pendingFlowId: null, updatedAt: new Date(), }) .where( and( eq(generations.pendingFlowId, flowId), eq(generations.projectId, projectId) ) ); // Also update associated images to have the flowId const generationIds = pendingGens.map(g => g.id); if (generationIds.length > 0) { await db .update(images) .set({ flowId: flowId, updatedAt: new Date(), }) .where( and( eq(images.projectId, projectId), isNull(images.flowId), inArray(images.generationId, generationIds) ) ); } } private async updateStatus( id: string, status: 'pending' | 'processing' | 'success' | 'failed', additionalUpdates?: { errorMessage?: string; outputImageId?: string; processingTimeMs?: number; } ): Promise { await db .update(generations) .set({ status, ...additionalUpdates, updatedAt: new Date(), }) .where(eq(generations.id, id)); } async getById(id: string): Promise { const generation = await db.query.generations.findFirst({ where: eq(generations.id, id), }); return generation || null; } async getByIdWithRelations(id: string): Promise { const generation = await db.query.generations.findFirst({ where: eq(generations.id, id), with: { outputImage: true, flow: true, }, }); if (!generation) { throw new Error(ERROR_MESSAGES.GENERATION_NOT_FOUND); } if (generation.referencedImages && Array.isArray(generation.referencedImages)) { const refImageIds = (generation.referencedImages as Array<{ imageId: string; alias: string }>) .map((ref) => ref.imageId); const refImages = await this.imageService.getMultipleByIds(refImageIds); return { ...generation, referenceImages: refImages, } as GenerationWithRelations; } return generation as GenerationWithRelations; } async list( filters: GenerationFilters, limit: number, offset: number ): Promise<{ generations: GenerationWithRelations[]; total: number }> { const conditions = [ buildEqCondition(generations, 'projectId', filters.projectId), buildEqCondition(generations, 'flowId', filters.flowId), buildEqCondition(generations, 'status', filters.status), ]; const whereClause = buildWhereClause(conditions); const [generationsList, countResult] = await Promise.all([ db.query.generations.findMany({ where: whereClause, orderBy: [desc(generations.createdAt)], limit, offset, with: { outputImage: true, flow: true, }, }), db .select({ count: count() }) .from(generations) .where(whereClause), ]); const totalCount = countResult[0]?.count || 0; return { generations: generationsList as GenerationWithRelations[], total: Number(totalCount), }; } /** * Regenerate an existing generation (Section 3) * - Allows regeneration for any status (no status checks) * - Uses exact same parameters as original * - Updates existing image (same ID, path, URL) * - No retry count logic */ async regenerate(id: string): Promise { const generation = await this.getById(id); if (!generation) { throw new Error(ERROR_MESSAGES.GENERATION_NOT_FOUND); } if (!generation.outputImageId) { throw new Error('Cannot regenerate generation without output image'); } const startTime = Date.now(); try { // Update status to processing await this.updateStatus(id, 'processing'); // Use EXACT same parameters as original (no overrides) const genResult = await this.imageGenService.generateImage({ prompt: generation.prompt, filename: `gen_${id}`, referenceImages: [], // TODO: Re-resolve referenced images if needed aspectRatio: generation.aspectRatio || GENERATION_LIMITS.DEFAULT_ASPECT_RATIO, orgId: 'default', projectId: generation.projectId, meta: generation.meta as Record || {}, }); if (!genResult.success) { const processingTime = Date.now() - startTime; await this.updateStatus(id, 'failed', { errorMessage: genResult.error || 'Regeneration failed', processingTimeMs: processingTime, }); throw new Error(genResult.error || 'Regeneration failed'); } // Note: Physical file in MinIO is overwritten by ImageGenService // Image record preserves: imageId, storageKey, storageUrl, alias, createdAt // Image record updates: fileSize (if changed), updatedAt const processingTime = Date.now() - startTime; await this.updateStatus(id, 'success', { processingTimeMs: processingTime, }); return await this.getByIdWithRelations(id); } catch (error) { const processingTime = Date.now() - startTime; await this.updateStatus(id, 'failed', { errorMessage: error instanceof Error ? error.message : 'Unknown error', processingTimeMs: processingTime, }); throw error; } } // Keep retry() for backward compatibility, delegate to regenerate() async retry(id: string, overrides?: { prompt?: string; aspectRatio?: string }): Promise { // Ignore overrides, regenerate with original parameters return await this.regenerate(id); } async update( id: string, updates: { prompt?: string; aspectRatio?: string; flowId?: string | null; meta?: Record; } ): Promise { const generation = await this.getById(id); if (!generation) { throw new Error(ERROR_MESSAGES.GENERATION_NOT_FOUND); } // Check if generative parameters changed (prompt or aspectRatio) const shouldRegenerate = (updates.prompt !== undefined && updates.prompt !== generation.prompt) || (updates.aspectRatio !== undefined && updates.aspectRatio !== generation.aspectRatio); // Handle flowId change (Section 9.2) if (updates.flowId !== undefined && updates.flowId !== null) { // If flowId provided and not null, create flow if it doesn't exist (eager creation) const existingFlow = await db.query.flows.findFirst({ where: eq(flows.id, updates.flowId), }); if (!existingFlow) { await db.insert(flows).values({ id: updates.flowId, projectId: generation.projectId, aliases: {}, meta: {}, }); } } // Update database fields const updateData: Partial = {}; if (updates.prompt !== undefined) { updateData.prompt = updates.prompt; // Update the prompt used for generation } if (updates.aspectRatio !== undefined) { updateData.aspectRatio = updates.aspectRatio; } if (updates.flowId !== undefined) { updateData.flowId = updates.flowId; } if (updates.meta !== undefined) { updateData.meta = updates.meta; } if (Object.keys(updateData).length > 0) { await db .update(generations) .set({ ...updateData, updatedAt: new Date() }) .where(eq(generations.id, id)); } // If generative parameters changed, trigger regeneration if (shouldRegenerate && generation.outputImageId) { // Update status to processing await this.updateStatus(id, 'processing'); try { // Use updated prompt/aspectRatio or fall back to existing const promptToUse = updates.prompt || generation.prompt; const aspectRatioToUse = updates.aspectRatio || generation.aspectRatio || GENERATION_LIMITS.DEFAULT_ASPECT_RATIO; // Regenerate image const genResult = await this.imageGenService.generateImage({ prompt: promptToUse, filename: `gen_${id}`, referenceImages: [], aspectRatio: aspectRatioToUse, orgId: 'default', projectId: generation.projectId, meta: updates.meta || generation.meta || {}, }); if (!genResult.success) { await this.updateStatus(id, 'failed', { errorMessage: genResult.error || 'Regeneration failed', }); throw new Error(genResult.error || 'Regeneration failed'); } // Note: Physical file in MinIO is overwritten by ImageGenService // TODO: Update fileSize and other metadata when ImageService.update() supports it await this.updateStatus(id, 'success'); } catch (error) { await this.updateStatus(id, 'failed', { errorMessage: error instanceof Error ? error.message : 'Unknown error', }); throw error; } } return await this.getByIdWithRelations(id); } /** * Conditional delete for generation (Section 7.2) * - If output image WITHOUT project alias → delete image + generation * - If output image WITH project alias → keep image, delete generation only, set generationId=NULL */ async delete(id: string): Promise { const generation = await this.getById(id); if (!generation) { throw new Error(ERROR_MESSAGES.GENERATION_NOT_FOUND); } if (generation.outputImageId) { // Get the output image to check if it has a project alias const outputImage = await this.imageService.getById(generation.outputImageId); if (outputImage) { if (outputImage.alias) { // Case 2: Image has project alias → keep image, delete generation only // Set generationId = NULL in image record await db .update(images) .set({ generationId: null, updatedAt: new Date() }) .where(eq(images.id, outputImage.id)); } else { // Case 1: Image has no alias → delete both image and generation await this.imageService.hardDelete(generation.outputImageId); } } } // Delete generation record (hard delete) await db.delete(generations).where(eq(generations.id, id)); } }