feature:FUR-31 [BE][Web] API check limit plan AI generation (#115)

src/ai-generation/contracts/constant.ts

@@ -26,4 +26,11 @@ export enum AIGenerationTaskProgress {
   QUEUEING = 0,
   RUNNING = 'RUNNING',
   SUCCESS = 100
+}
+
+export const DEFAULT_CREDITS = 50;
+
+export enum AIGenerationPricing {
+  TEXT_TO_MODEL = 15,
+  TEXT_TO_IMAGE = 10
 }

src/ai-generation/services/text-to-image.service.ts

@@ -6,8 +6,10 @@ import { catchError, firstValueFrom } from 'rxjs'
 import { AxiosError } from 'axios'
 import { AppException } from '@common/exceptions/app.exception'
 import { Errors } from '@common/contracts/error'
-import { AIGenerationPlatform, AIGenerationType } from '@ai-generation/contracts/constant'
+import { AIGenerationPlatform, AIGenerationPricing, AIGenerationType } from '@ai-generation/contracts/constant'
 import { GenerateTextToImageDto } from '@ai-generation/dtos/text-to-image.dto'
+import { CustomerRepository } from '@customer/repositories/customer.repository'
+import { Status } from '@common/contracts/constant'
 
 @Injectable()
 export class AIGenerationTextToImageService {
@@ -17,7 +19,8 @@ export class AIGenerationTextToImageService {
   constructor(
     private readonly aiGenerationRepository: AIGenerationRepository,
     private readonly httpService: HttpService,
-    private readonly configService: ConfigService
+    private readonly configService: ConfigService,
+    private readonly customerRepository: CustomerRepository
   ) {
     this.config = this.configService.get('edenAI')
     this.headersRequest = {
@@ -29,7 +32,16 @@ export class AIGenerationTextToImageService {
   async generateTextToImage(generateTextToImageDto: GenerateTextToImageDto) {
     const { customerId } = generateTextToImageDto
 
-    // TODO: Check limit AI generation here
+    // Check limit AI generation
+    const { credits } = await this.customerRepository.findOne({
+      conditions: {
+        _id: customerId,
+        status: Status.ACTIVE
+      }
+    })
+    if (credits < AIGenerationPricing.TEXT_TO_IMAGE) {
+      throw new AppException(Errors.NOT_ENOUGH_CREDITS_ERROR)
+    }
 
     const { data } = await firstValueFrom(
       this.httpService
@@ -47,14 +59,21 @@ export class AIGenerationTextToImageService {
     if (result?.status !== 'success') throw new AppException({ ...Errors.EDEN_AI_ERROR, data })
 
     const imageUrl = result?.items[0]?.image_resource_url
-    await this.aiGenerationRepository.create({
-      customerId,
-      type: AIGenerationType.TEXT_TO_IMAGE,
-      platform: AIGenerationPlatform.EDEN_AI,
-      cost: result?.cost ?? 0.01, // total 1 credits
-      imageUrl
-    })
-
+    await Promise.all([
+      this.aiGenerationRepository.create({
+        customerId,
+        type: AIGenerationType.TEXT_TO_IMAGE,
+        platform: AIGenerationPlatform.EDEN_AI,
+        cost: result?.cost ?? 0.01, // total 1 credits
+        imageUrl
+      }),
+      this.customerRepository.findOneAndUpdate(
+        { _id: customerId },
+        {
+          $inc: { credits: -AIGenerationPricing.TEXT_TO_IMAGE }
+        }
+      )
+    ])
     return { imageUrl }
   }
 }

src/ai-generation/services/text-to-model.service.ts

@@ -7,7 +7,9 @@ import { AxiosError } from 'axios'
 import { GenerateTextToDraftModelDto } from '@ai-generation/dtos/text-to-model.dto'
 import { AppException } from '@common/exceptions/app.exception'
 import { Errors } from '@common/contracts/error'
-import { AIGenerationPlatform, AIGenerationType } from '@ai-generation/contracts/constant'
+import { AIGenerationPlatform, AIGenerationPricing, AIGenerationType } from '@ai-generation/contracts/constant'
+import { CustomerRepository } from '@customer/repositories/customer.repository'
+import { Status } from '@common/contracts/constant'
 
 @Injectable()
 export class AIGenerationTextToModelService {
@@ -18,6 +20,7 @@ export class AIGenerationTextToModelService {
     private readonly aiGenerationRepository: AIGenerationRepository,
     private readonly httpService: HttpService,
     private readonly configService: ConfigService,
+    private readonly customerRepository: CustomerRepository
   ) {
     this.config = this.configService.get('tripo3dAI')
     this.headersRequest = {
@@ -27,9 +30,18 @@ export class AIGenerationTextToModelService {
   }
 
   async generateTextToDraftModel(generateTextToDraftModelDto: GenerateTextToDraftModelDto) {
-    const {customerId} = generateTextToDraftModelDto
+    const { customerId } = generateTextToDraftModelDto
 
-    // TODO: Check limit AI generation here
+    // Check limit AI generation
+    const { credits } = await this.customerRepository.findOne({
+      conditions: {
+        _id: customerId,
+        status: Status.ACTIVE
+      }
+    })
+    if (credits < AIGenerationPricing.TEXT_TO_MODEL) {
+      throw new AppException(Errors.NOT_ENOUGH_CREDITS_ERROR)
+    }
 
     const { data } = await firstValueFrom(
       this.httpService
@@ -45,13 +57,21 @@ export class AIGenerationTextToModelService {
     )
     if (data.code !== 0) throw new AppException({ ...Errors.TRIPO_3D_AI_ERROR, data })
 
-    await this.aiGenerationRepository.create({
-      customerId,
-      type: AIGenerationType.TEXT_TO_MODEL,
-      platform: AIGenerationPlatform.TRIPO_3D_AI,
-      cost: 20, // total 2000 credits
-      taskId: data?.data?.task_id
-    })
+    await Promise.all([
+      this.aiGenerationRepository.create({
+        customerId,
+        type: AIGenerationType.TEXT_TO_MODEL,
+        platform: AIGenerationPlatform.TRIPO_3D_AI,
+        cost: 20, // total 2000 credits
+        taskId: data?.data?.task_id
+      }),
+      this.customerRepository.findOneAndUpdate(
+        { _id: customerId },
+        {
+          $inc: { credits: -AIGenerationPricing.TEXT_TO_MODEL }
+        }
+      )
+    ])
 
     return data?.data
   }

src/common/contracts/error.ts

@@ -125,5 +125,10 @@ export const Errors: Record<string, ErrorResponse> = {
     error: 'EDEN_AI_ERROR',
     message: 'Có chút lỗi xảy ra. Vui lòng thử lại sau giây lát bạn nhé.',
     httpStatus: HttpStatus.INTERNAL_SERVER_ERROR
+  },
+  NOT_ENOUGH_CREDITS_ERROR: {
+    error: 'NOT_ENOUGH_CREDITS_ERROR',
+    message: 'Số lượng credits còn lại không đủ. Vui lòng nạp thêm nhé.',
+    httpStatus: HttpStatus.BAD_REQUEST
   }
 }

src/customer/customer.module.ts

@@ -1,14 +1,29 @@
-import { Module } from '@nestjs/common'
+import { Module, OnModuleInit } from '@nestjs/common'
 import { MongooseModule } from '@nestjs/mongoose'
 import { Customer, CustomerSchema } from '@customer/schemas/customer.schema'
 import { CustomerRepository } from '@customer/repositories/customer.repository'
 import { CustomerService } from '@customer/services/customer.service'
 import { CustomerController } from '@src/customer/controllers/customer.controller'
+import { DEFAULT_CREDITS } from '@ai-generation/contracts/constant'
 
 @Module({
   imports: [MongooseModule.forFeature([{ name: Customer.name, schema: CustomerSchema }])],
   controllers: [CustomerController],
   providers: [CustomerService, CustomerRepository],
-  exports: [CustomerService, CustomerRepository],
+  exports: [CustomerService, CustomerRepository]
 })
-export class CustomerModule {}
+export class CustomerModule implements OnModuleInit {
+  constructor(private readonly customerRepository: CustomerRepository) {}
+
+  async onModuleInit() {
+    console.log(`CustomerModule.OnModuleInit: Set default credits for customers`)
+    await this.customerRepository.updateMany(
+      { credits: { $exists: false } },
+      {
+        $set: {
+          credits: DEFAULT_CREDITS
+        }
+      }
+    )
+  }
+}

src/customer/schemas/customer.schema.ts

@@ -4,6 +4,7 @@ import * as paginate from 'mongoose-paginate-v2'
 import { ApiProperty } from '@nestjs/swagger'
 import { Transform } from 'class-transformer'
 import { Gender, Status } from '@common/contracts/constant'
+import { DEFAULT_CREDITS } from '@ai-generation/contracts/constant'
 
 export type CustomerDocument = HydratedDocument<Customer>
 
@@ -76,6 +77,10 @@ export class Customer {
   @Prop({ type: String })
   googleUserId: string
 
+  @ApiProperty()
+  @Prop({ type: Number, min: 0, default: DEFAULT_CREDITS })
+  credits: number
+
   @Prop({
     enum: Status,
     default: Status.ACTIVE