diff --git a/src/quota/tokenCount.ts b/src/quota/tokenCount.ts index 4c24fbb..70ab318 100644 --- a/src/quota/tokenCount.ts +++ b/src/quota/tokenCount.ts @@ -12,16 +12,19 @@ import { Usage } from "@prisma/client"; export default class tokenCount implements IQuota { defaultQuota: number; lookback: number; - considerInputTokensAsHalf: boolean; - + requestTokenMultiplier: number; + responseTokenMultiplier: number; + constructor( defaultQuota: number = 512 * 25, lookback: number = 1000 * 60 * 60 * 24, - considerInputTokensAsHalf: boolean = true, + requestTokenMultiplier: number = 1, + responseTokenMultiplier: number = 1, ) { this.defaultQuota = defaultQuota; this.lookback = lookback; - this.considerInputTokensAsHalf = considerInputTokensAsHalf; + this.requestTokenMultiplier = requestTokenMultiplier; + this.responseTokenMultiplier = responseTokenMultiplier; } private getUserQuota(id: string) { @@ -55,9 +58,7 @@ export default class tokenCount implements IQuota { const usageResponse = usedTokens.usageResponse === null ? 0 : usedTokens.usageResponse; const usedUnits = (() => { - if (this.considerInputTokensAsHalf) - return usageResponse + usageRequest / 2; - return usageResponse + usageRequest; + return usageRequest * this.requestTokenMultiplier + usageResponse * this.responseTokenMultiplier; })(); if (userQuota?.vip) return this.createUserQuotaData(Infinity, usedUnits); @@ -74,30 +75,13 @@ export default class tokenCount implements IQuota { * @returns promise of giving out the record */ findNthUsage(user: string, requestTimestamp: number, unitCount: number) { - if (this.considerInputTokensAsHalf) - return database.$queryRaw>` - SELECT t1.*, ( - SELECT - SUM(usageResponse + usageRequest/2) AS usage - FROM \`usage\` - WHERE - user = ${user} AND - timestamp >= ${requestTimestamp - this.lookback} AND - timestamp <= t1.timestamp - ) as usage - FROM - \`usage\` AS t1 - WHERE - user = ${user} AND - timestamp >= ${requestTimestamp - this.lookback} AND - usage >= ${unitCount} - ORDER BY timestamp ASC - LIMIT 1 - `; - return database.$queryRaw>` + return database.$queryRaw>` SELECT t1.*, ( SELECT - SUM(usageResponse + usageRequest) AS usage + SUM( + usageRequest * ${this.requestTokenMultiplier} + + usageResponse * ${this.responseTokenMultiplier} + ) AS usage FROM \`usage\` WHERE user = ${user} AND