diff --git a/src/ai-generation/contracts/constant.ts b/src/ai-generation/contracts/constant.ts index 8913af8..34e5989 100644 --- a/src/ai-generation/contracts/constant.ts +++ b/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 } \ No newline at end of file diff --git a/src/ai-generation/services/text-to-image.service.ts b/src/ai-generation/services/text-to-image.service.ts index d2d4b06..5efe239 100644 --- a/src/ai-generation/services/text-to-image.service.ts +++ b/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 } } } diff --git a/src/ai-generation/services/text-to-model.service.ts b/src/ai-generation/services/text-to-model.service.ts index 616c42a..ece542a 100644 --- a/src/ai-generation/services/text-to-model.service.ts +++ b/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 } diff --git a/src/common/contracts/error.ts b/src/common/contracts/error.ts index 8584581..ffadece 100644 --- a/src/common/contracts/error.ts +++ b/src/common/contracts/error.ts @@ -125,5 +125,10 @@ export const Errors: Record = { 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 } } diff --git a/src/customer/customer.module.ts b/src/customer/customer.module.ts index c16f751..40cd320 100644 --- a/src/customer/customer.module.ts +++ b/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 + } + } + ) + } +} diff --git a/src/customer/schemas/customer.schema.ts b/src/customer/schemas/customer.schema.ts index 568d3d4..ca8961e 100644 --- a/src/customer/schemas/customer.schema.ts +++ b/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 @@ -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