diff --git a/src/modules/ai-runtime/ai-runtime.module.ts b/src/modules/ai-runtime/ai-runtime.module.ts index 18255bc..c0f8ce9 100644 --- a/src/modules/ai-runtime/ai-runtime.module.ts +++ b/src/modules/ai-runtime/ai-runtime.module.ts @@ -6,11 +6,12 @@ import { UserAiService } from './user-ai.service'; import { CredentialEncryptionService } from './credential-encryption.service'; import { RuntimeInternalController } from './internal/runtime-internal.controller'; import { RuntimeInternalService } from './internal/runtime-internal.service'; +import { UserAiQuotaService } from './user-ai-quota.service'; @Module({ imports: [ConfigModule, PrismaModule], controllers: [UserAiController, RuntimeInternalController], - providers: [UserAiService, CredentialEncryptionService, RuntimeInternalService], - exports: [UserAiService, CredentialEncryptionService, RuntimeInternalService], + providers: [UserAiService, CredentialEncryptionService, RuntimeInternalService, UserAiQuotaService], + exports: [UserAiService, CredentialEncryptionService, RuntimeInternalService, UserAiQuotaService], }) export class AiRuntimeModule {} diff --git a/src/modules/ai-runtime/user-ai-quota.service.ts b/src/modules/ai-runtime/user-ai-quota.service.ts new file mode 100644 index 0000000..f2287df --- /dev/null +++ b/src/modules/ai-runtime/user-ai-quota.service.ts @@ -0,0 +1,97 @@ +import { Injectable, BadRequestException } from '@nestjs/common'; +import { PrismaService } from '../../infrastructure/database/prisma.service'; + +@Injectable() +export class UserAiQuotaService { + constructor(private readonly prisma: PrismaService) {} + + /** Get today's date truncated to midnight UTC */ + private today(): Date { + const d = new Date(); + return new Date(Date.UTC(d.getFullYear(), d.getMonth(), d.getDate())); + } + + /** Check if user can create a new AI job, throw if over limit */ + async checkAndReserve(userId: string, apiKeyMode: string): Promise { + const settings = await this.prisma.userAiSettings.findUnique({ where: { userId } }); + const maxJobs = settings?.maxDailyAiJobs ?? 20; + const maxTokens = settings?.maxDailyTokenBudget ?? 100_000; + + const localDate = this.today(); + + const usage = await this.prisma.userAiUsageDaily.findUnique({ + where: { userId_localDate_apiKeyMode: { userId, localDate, apiKeyMode } }, + }); + + const jobCount = usage?.jobCount ?? 0; + const totalTokens = usage?.totalTokens ?? 0; + + if (jobCount >= maxJobs) { + throw new BadRequestException({ + errorCode: 'DAILY_JOB_LIMIT_EXCEEDED', + message: `Daily AI job limit (${maxJobs}) exceeded. Current: ${jobCount}`, + }); + } + + if (totalTokens >= maxTokens) { + throw new BadRequestException({ + errorCode: 'DAILY_TOKEN_BUDGET_EXCEEDED', + message: `Daily token budget (${maxTokens}) exceeded. Current: ${totalTokens}`, + }); + } + } + + /** Increment job count (called on job creation) */ + async incrementJobCount(userId: string, apiKeyMode: string): Promise { + const localDate = this.today(); + await this.prisma.userAiUsageDaily.upsert({ + where: { userId_localDate_apiKeyMode: { userId, localDate, apiKeyMode } }, + create: { userId, localDate, apiKeyMode, jobCount: 1 }, + update: { jobCount: { increment: 1 } }, + }); + } + + /** Record token consumption from invocation log (called after Runtime call) */ + async recordTokenUsage( + userId: string, + apiKeyMode: string, + inputTokens: number, + outputTokens: number, + totalTokens: number, + costEstimate?: number, + ): Promise { + const localDate = this.today(); + await this.prisma.userAiUsageDaily.upsert({ + where: { userId_localDate_apiKeyMode: { userId, localDate, apiKeyMode } }, + create: { + userId, localDate, apiKeyMode, + inputTokens, outputTokens, totalTokens, + costEstimate: costEstimate ?? 0, + jobCount: 0, + }, + update: { + inputTokens: { increment: inputTokens }, + outputTokens: { increment: outputTokens }, + totalTokens: { increment: totalTokens }, + costEstimate: { increment: costEstimate ?? 0 }, + }, + }); + } + + /** Get today's usage summary for a user */ + async getUsage(userId: string) { + const localDate = this.today(); + const [platform, userKey] = await Promise.all([ + this.prisma.userAiUsageDaily.findUnique({ + where: { userId_localDate_apiKeyMode: { userId, localDate, apiKeyMode: 'platform_key' } }, + }), + this.prisma.userAiUsageDaily.findUnique({ + where: { userId_localDate_apiKeyMode: { userId, localDate, apiKeyMode: 'user_deepseek_key' } }, + }), + ]); + return { + platformKey: platform ?? { jobCount: 0, totalTokens: 0 }, + userDeepseekKey: userKey ?? { jobCount: 0, totalTokens: 0 }, + }; + } +}