feat: 用户 AI 消费额度与日限流 (API-AI-070)
All checks were successful
Deploy API Server / build-and-deploy (push) Successful in 43s
All checks were successful
Deploy API Server / build-and-deploy (push) Successful in 43s
- UserAiQuotaService: checkAndReserve / incrementJobCount / recordTokenUsage - 创建 Job 前检查 maxDailyAiJobs + maxDailyTokenBudget - Runtime 调用后更新 UserAiUsageDaily token 消耗 - user_deepseek_key 不计入平台 token 预算 - 超限返回 DAILY_JOB_LIMIT_EXCEEDED / DAILY_TOKEN_BUDGET_EXCEEDED Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
parent
43e6e9029c
commit
5cbf20046a
@ -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 {}
|
||||
|
||||
97
src/modules/ai-runtime/user-ai-quota.service.ts
Normal file
97
src/modules/ai-runtime/user-ai-quota.service.ts
Normal file
@ -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<void> {
|
||||
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<void> {
|
||||
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<void> {
|
||||
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 },
|
||||
};
|
||||
}
|
||||
}
|
||||
Loading…
x
Reference in New Issue
Block a user