@lobehub/chat 1.98.2 → 1.99.1

package/src/database/models/generation.ts

@@ -0,0 +1,197 @@
+import debug from 'debug';
+import { and, eq } from 'drizzle-orm';
+
+import { LobeChatDatabase, Transaction } from '@/database/type';
+import { FileService } from '@/server/services/file';
+import { AsyncTaskError, AsyncTaskStatus } from '@/types/asyncTask';
+import { FileSource } from '@/types/files';
+import { Generation, ImageGenerationAsset } from '@/types/generation';
+
+import { NewFile } from '../schemas';
+import {
+  GenerationItem,
+  GenerationWithAsyncTask,
+  NewGeneration,
+  generations,
+} from '../schemas/generation';
+import { FileModel } from './file';
+
+// Create debug logger
+const log = debug('lobe-image:generation-model');
+
+export class GenerationModel {
+  private db: LobeChatDatabase;
+  private userId: string;
+  private fileModel: FileModel;
+  private fileService: FileService;
+
+  constructor(db: LobeChatDatabase, userId: string) {
+    this.db = db;
+    this.userId = userId;
+    this.fileModel = new FileModel(db, userId);
+    this.fileService = new FileService(db, userId);
+  }
+
+  async create(value: Omit<NewGeneration, 'userId'>): Promise<GenerationItem> {
+    log('Creating generation: %O', {
+      generationBatchId: value.generationBatchId,
+      userId: this.userId,
+    });
+
+    const [result] = await this.db
+      .insert(generations)
+      .values({ ...value, userId: this.userId })
+      .returning();
+
+    log('Generation created successfully: %s', result.id);
+    return result;
+  }
+
+  async findById(id: string): Promise<GenerationItem | undefined> {
+    log('Finding generation by ID: %s for user: %s', id, this.userId);
+
+    const result = await this.db.query.generations.findFirst({
+      where: and(eq(generations.id, id), eq(generations.userId, this.userId)),
+    });
+
+    log('Generation %s: %s', id, result ? 'found' : 'not found');
+    return result;
+  }
+
+  async findByIdWithAsyncTask(id: string): Promise<GenerationWithAsyncTask | undefined> {
+    log('Finding generation by ID: %s for user: %s', id, this.userId);
+
+    const result = await this.db.query.generations.findFirst({
+      where: and(eq(generations.id, id), eq(generations.userId, this.userId)),
+      with: {
+        asyncTask: true,
+      },
+    });
+
+    log('Generation %s: %s', id, result ? 'found' : 'not found');
+    return result as GenerationWithAsyncTask | undefined;
+  }
+
+  async update(id: string, value: Partial<NewGeneration>, trx?: Transaction) {
+    log('Updating generation: %s with values: %O', id, {
+      asyncTaskId: value.asyncTaskId,
+      hasAsset: !!value.asset,
+    });
+
+    const executeUpdate = async (tx: Transaction) => {
+      return await tx
+        .update(generations)
+        .set({ ...value, updatedAt: new Date() })
+        .where(and(eq(generations.id, id), eq(generations.userId, this.userId)));
+    };
+
+    const result = await (trx ? executeUpdate(trx) : this.db.transaction(executeUpdate));
+
+    log('Generation %s updated successfully', id);
+    return result;
+  }
+
+  async createAssetAndFile(
+    id: string,
+    asset: ImageGenerationAsset,
+    file: Omit<NewFile, 'id' | 'userId'>,
+  ) {
+    log('Creating generation asset and file with transaction: %s', id);
+
+    return await this.db.transaction(async (tx: Transaction) => {
+      // Create file first using transaction
+      // Since duplicates are very rare, we always create globalFile - checking existence first would be wasteful
+      const newFile = await this.fileModel.create(
+        {
+          ...file,
+          source: FileSource.ImageGeneration,
+        },
+        true,
+        tx,
+      );
+
+      // Update generation with asset and fileId using the transaction-aware update method
+      await this.update(
+        id,
+        {
+          asset,
+          fileId: newFile.id,
+        },
+        tx,
+      );
+
+      log('Generation %s updated with asset and file %s successfully', id, newFile.id);
+
+      return {
+        file: newFile,
+      };
+    });
+  }
+
+  async delete(id: string, trx?: Transaction) {
+    log('Deleting generation: %s for user: %s', id, this.userId);
+
+    const executeDelete = async (tx: Transaction) => {
+      return await tx
+        .delete(generations)
+        .where(and(eq(generations.id, id), eq(generations.userId, this.userId)))
+        .returning();
+    };
+
+    const result = await (trx ? executeDelete(trx) : this.db.transaction(executeDelete));
+    const deletedGeneration = result[0];
+
+    log('Generation %s deleted successfully', id);
+    return deletedGeneration;
+  }
+
+  /**
+   * Find generation by ID and transform it to frontend type
+   * This method uses findByIdWithAsyncTask and applies transformation
+   */
+  async findByIdAndTransform(id: string): Promise<Generation | null> {
+    log('Finding and transforming generation: %s', id);
+
+    const generation = await this.findByIdWithAsyncTask(id);
+    if (!generation) {
+      log('Generation %s not found', id);
+      return null;
+    }
+
+    return await this.transformGeneration(generation);
+  }
+
+  /**
+   * Transform a GenerationItem (database type) to Generation (frontend type)
+   * This method processes asset URLs and async task information
+   */
+  async transformGeneration(generation: GenerationWithAsyncTask): Promise<Generation> {
+    // Process asset URLs if they exist, following the same logic as in generationBatch.ts
+    const asset = generation.asset as ImageGenerationAsset | null;
+    if (asset && asset.url && asset.thumbnailUrl) {
+      const [url, thumbnailUrl] = await Promise.all([
+        this.fileService.getFullFileUrl(asset.url),
+        this.fileService.getFullFileUrl(asset.thumbnailUrl),
+      ]);
+      asset.url = url;
+      asset.thumbnailUrl = thumbnailUrl;
+    }
+
+    // Build the Generation object following the same structure as in generationBatch.ts
+    const result: Generation = {
+      asset,
+      asyncTaskId: generation.asyncTaskId || null,
+      createdAt: generation.createdAt,
+      id: generation.id,
+      seed: generation.seed,
+      task: {
+        error: generation.asyncTask?.error
+          ? (generation.asyncTask.error as AsyncTaskError)
+          : undefined,
+        id: generation.asyncTaskId || '',
+        status: (generation.asyncTask?.status as AsyncTaskStatus) || 'pending',
+      },
+    };
+    return result;
+  }
+}

package/src/database/models/generationBatch.ts

@@ -0,0 +1,212 @@
+import debug from 'debug';
+import { and, eq } from 'drizzle-orm';
+
+import { LobeChatDatabase } from '@/database/type';
+import { FileService } from '@/server/services/file';
+import { Generation, GenerationAsset, GenerationBatch, GenerationConfig } from '@/types/generation';
+
+import {
+  GenerationBatchItem,
+  GenerationBatchWithGenerations,
+  NewGenerationBatch,
+  generationBatches,
+} from '../schemas/generation';
+import { GenerationModel } from './generation';
+
+const log = debug('lobe-image:generation-batch-model');
+
+export class GenerationBatchModel {
+  private db: LobeChatDatabase;
+  private userId: string;
+  private fileService: FileService;
+  private generationModel: GenerationModel;
+
+  constructor(db: LobeChatDatabase, userId: string) {
+    this.db = db;
+    this.userId = userId;
+    this.fileService = new FileService(db, userId);
+    this.generationModel = new GenerationModel(db, userId);
+  }
+
+  async create(value: NewGenerationBatch): Promise<GenerationBatchItem> {
+    log('Creating generation batch: %O', {
+      topicId: value.generationTopicId,
+      userId: this.userId,
+    });
+
+    const [result] = await this.db
+      .insert(generationBatches)
+      .values({ ...value, userId: this.userId })
+      .returning();
+
+    log('Generation batch created successfully: %s', result.id);
+    return result;
+  }
+
+  async findById(id: string): Promise<GenerationBatchItem | undefined> {
+    log('Finding generation batch by ID: %s for user: %s', id, this.userId);
+
+    const result = await this.db.query.generationBatches.findFirst({
+      where: and(eq(generationBatches.id, id), eq(generationBatches.userId, this.userId)),
+    });
+
+    log('Generation batch %s: %s', id, result ? 'found' : 'not found');
+    return result;
+  }
+
+  async findByTopicId(topicId: string): Promise<GenerationBatchItem[]> {
+    log('Finding generation batches by topic ID: %s for user: %s', topicId, this.userId);
+
+    const results = await this.db.query.generationBatches.findMany({
+      orderBy: (table, { desc }) => [desc(table.createdAt)],
+      where: and(
+        eq(generationBatches.generationTopicId, topicId),
+        eq(generationBatches.userId, this.userId),
+      ),
+    });
+
+    log('Found %d generation batches for topic %s', results.length, topicId);
+    return results;
+  }
+
+  /**
+   * Find batches with their associated generations using relations
+   */
+  async findByTopicIdWithGenerations(topicId: string): Promise<GenerationBatchWithGenerations[]> {
+    log(
+      'Finding generation batches with generations for topic ID: %s for user: %s',
+      topicId,
+      this.userId,
+    );
+
+    const results = await this.db.query.generationBatches.findMany({
+      orderBy: (table, { asc }) => [asc(table.createdAt)],
+      where: and(
+        eq(generationBatches.generationTopicId, topicId),
+        eq(generationBatches.userId, this.userId),
+      ),
+      with: {
+        generations: {
+          orderBy: (table, { asc }) => [asc(table.createdAt), asc(table.id)],
+          with: {
+            asyncTask: true,
+          },
+        },
+      },
+    });
+
+    log('Found %d generation batches with generations for topic %s', results.length, topicId);
+    return results as GenerationBatchWithGenerations[];
+  }
+
+  async queryGenerationBatchesByTopicIdWithGenerations(
+    topicId: string,
+  ): Promise<(GenerationBatch & { generations: Generation[] })[]> {
+    log('Fetching generation batches for topic ID: %s for user: %s', topicId, this.userId);
+
+    const batchesWithGenerations = await this.findByTopicIdWithGenerations(topicId);
+    if (batchesWithGenerations.length === 0) {
+      log('No batches found for topic: %s', topicId);
+      return [];
+    }
+
+    // Transform the database result to match our frontend types
+    const result: GenerationBatch[] = await Promise.all(
+      batchesWithGenerations.map(async (batch) => {
+        const [generations, config] = await Promise.all([
+          // Transform generations
+          Promise.all(
+            batch.generations.map((gen) => this.generationModel.transformGeneration(gen)),
+          ),
+          // Transform config
+          (async () => {
+            const config = batch.config as GenerationConfig;
+            if (Array.isArray(config.imageUrls)) {
+              config.imageUrls = await Promise.all(
+                config.imageUrls.map((url) => this.fileService.getFullFileUrl(url)),
+              );
+            }
+            return config;
+          })(),
+        ]);
+
+        return {
+          config,
+          createdAt: batch.createdAt,
+          generations,
+          height: batch.height,
+          id: batch.id,
+          model: batch.model,
+          prompt: batch.prompt,
+          provider: batch.provider,
+          width: batch.width,
+        };
+      }),
+    );
+
+    log('Feed construction complete for topic: %s, returning %d batches', topicId, result.length);
+    return result;
+  }
+
+  /**
+   * Delete a generation batch and return associated file URLs for cleanup
+   *
+   * This method follows the "database first, files second" deletion principle:
+   * 1. First queries the batch with its generations to collect thumbnail URLs
+   * 2. Then deletes the database record (cascade delete handles related generations)
+   * 3. Returns the deleted batch data and thumbnail URLs for file cleanup
+   *
+   * @param id - The batch ID to delete
+   * @returns Object containing deleted batch data and thumbnail URLs to clean, or undefined if batch not found or access denied
+   */
+  async delete(
+    id: string,
+  ): Promise<{ deletedBatch: GenerationBatchItem; thumbnailUrls: string[] } | undefined> {
+    log('Deleting generation batch: %s for user: %s', id, this.userId);
+
+    // 1. First, get generations with their assets to collect thumbnail URLs
+    const batchWithGenerations = await this.db.query.generationBatches.findFirst({
+      where: and(eq(generationBatches.id, id), eq(generationBatches.userId, this.userId)),
+      with: {
+        generations: {
+          columns: {
+            asset: true,
+          },
+        },
+      },
+    });
+
+    // If batch doesn't exist or doesn't belong to user, return undefined
+    if (!batchWithGenerations) {
+      return undefined;
+    }
+
+    // 2. Collect thumbnail URLs that need to be deleted
+    const thumbnailUrls: string[] = [];
+    if (batchWithGenerations.generations) {
+      for (const gen of batchWithGenerations.generations) {
+        const asset = gen.asset as GenerationAsset;
+        if (asset?.thumbnailUrl) {
+          thumbnailUrls.push(asset.thumbnailUrl);
+        }
+      }
+    }
+
+    // 3. Delete the batch record (this will cascade delete all associated generations)
+    const [deletedBatch] = await this.db
+      .delete(generationBatches)
+      .where(and(eq(generationBatches.id, id), eq(generationBatches.userId, this.userId)))
+      .returning();
+
+    log(
+      'Generation batch %s deleted successfully with %d thumbnails to clean',
+      id,
+      thumbnailUrls.length,
+    );
+
+    return {
+      deletedBatch,
+      thumbnailUrls,
+    };
+  }
+}

package/src/database/models/generationTopic.ts

@@ -0,0 +1,131 @@
+import { and, desc, eq } from 'drizzle-orm/expressions';
+
+import { LobeChatDatabase } from '@/database/type';
+import { FileService } from '@/server/services/file';
+import { GenerationAsset, ImageGenerationTopic } from '@/types/generation';
+
+import { GenerationTopicItem, generationTopics } from '../schemas/generation';
+
+export class GenerationTopicModel {
+  private userId: string;
+  private db: LobeChatDatabase;
+  private fileService: FileService;
+
+  constructor(db: LobeChatDatabase, userId: string) {
+    this.userId = userId;
+    this.db = db;
+    this.fileService = new FileService(db, userId);
+  }
+
+  queryAll = async () => {
+    const topics = await this.db
+      .select()
+      .from(generationTopics)
+      .orderBy(desc(generationTopics.updatedAt))
+      .where(eq(generationTopics.userId, this.userId));
+
+    return Promise.all(
+      topics.map(async (topic) => {
+        if (topic.coverUrl) {
+          return {
+            ...topic,
+            coverUrl: await this.fileService.getFullFileUrl(topic.coverUrl),
+          };
+        }
+        return topic;
+      }),
+    );
+  };
+
+  create = async (title: string) => {
+    const [newGenerationTopic] = await this.db
+      .insert(generationTopics)
+      .values({
+        title,
+        userId: this.userId,
+      })
+      .returning();
+
+    return newGenerationTopic;
+  };
+
+  update = async (
+    id: string,
+    data: Partial<ImageGenerationTopic>,
+  ): Promise<GenerationTopicItem | undefined> => {
+    const [updatedTopic] = await this.db
+      .update(generationTopics)
+      .set({ ...data, updatedAt: new Date() })
+      .where(and(eq(generationTopics.id, id), eq(generationTopics.userId, this.userId)))
+      .returning();
+
+    return updatedTopic;
+  };
+
+  /**
+   * Delete a generation topic and return associated file URLs for cleanup
+   *
+   * This method follows the "database first, files second" deletion principle:
+   * 1. First queries the topic with all its batches and generations to collect file URLs
+   * 2. Then deletes the database record (cascade delete handles related batches and generations)
+   * 3. Returns the deleted topic data and file URLs for cleanup
+   *
+   * @param id - The topic ID to delete
+   * @returns Object containing deleted topic data and file URLs to clean, or undefined if topic not found or access denied
+   */
+  delete = async (
+    id: string,
+  ): Promise<{ deletedTopic: GenerationTopicItem; filesToDelete: string[] } | undefined> => {
+    // 1. First, get the topic with all its batches and generations to collect file URLs
+    const topicWithBatches = await this.db.query.generationTopics.findFirst({
+      where: and(eq(generationTopics.id, id), eq(generationTopics.userId, this.userId)),
+      with: {
+        batches: {
+          with: {
+            generations: {
+              columns: {
+                asset: true,
+              },
+            },
+          },
+        },
+      },
+    });
+
+    // If topic doesn't exist or doesn't belong to user, return undefined
+    if (!topicWithBatches) {
+      return undefined;
+    }
+
+    // 2. Collect all file URLs that need to be deleted
+    const filesToDelete: string[] = [];
+
+    // Add cover image URL if exists
+    if (topicWithBatches.coverUrl) {
+      filesToDelete.push(topicWithBatches.coverUrl);
+    }
+
+    // Add thumbnail URLs from all generations
+    if (topicWithBatches.batches) {
+      for (const batch of topicWithBatches.batches) {
+        for (const gen of batch.generations) {
+          const asset = gen.asset as GenerationAsset;
+          if (asset?.thumbnailUrl) {
+            filesToDelete.push(asset.thumbnailUrl);
+          }
+        }
+      }
+    }
+
+    // 3. Delete the topic record (this will cascade delete all batches and generations)
+    const [deletedTopic] = await this.db
+      .delete(generationTopics)
+      .where(and(eq(generationTopics.id, id), eq(generationTopics.userId, this.userId)))
+      .returning();
+
+    return {
+      deletedTopic,
+      filesToDelete,
+    };
+  };
+}