diff --git a/src/commands/check-limit.ts b/src/commands/check-limit.ts index 616bb11..c9e667b 100644 --- a/src/commands/check-limit.ts +++ b/src/commands/check-limit.ts @@ -35,7 +35,7 @@ export default class MyLimit extends Command implements Command { fields.push({ name: "Usage", inline: true, value: `${userQuotaRecovery.used} ${userQuotaRecovery.unitName}`.trim() }); if (userQuotaRecovery.recoveryTimestamp !== undefined) fields.push({ - name: `Recovery for ${recoveryFor} message${recoveryFor>1 ? "s" : ""}`, + name: `Recovery for ${recoveryFor} ${userQuotaRecovery.unitName}`.trim(), value: userQuotaRecovery.recoveryTimestamp === Infinity ? "never" : `` }); diff --git a/src/quota/messageCount.ts b/src/quota/messageCount.ts index 2e9bc78..7b97755 100644 --- a/src/quota/messageCount.ts +++ b/src/quota/messageCount.ts @@ -11,6 +11,10 @@ export default class MessageCount implements IQuota { lookback: number; defaultQuota: number; + /** + * @param defaultQuota the default quota for users that don't have override + * @param lookback lookback to check + */ constructor( defaultQuota: number = 25, lookback: number = 1000 * 60 * 60 * 24 @@ -19,6 +23,9 @@ export default class MessageCount implements IQuota { this.lookback = lookback; } + /** + * Retrives the quota from the database + */ private getUserQuota(id: string) { return database.limits.findUnique({ where: { user: BigInt(id) }, @@ -78,6 +85,9 @@ export default class MessageCount implements IQuota { }; } + /** + * helper funtion to create userQuotaData + */ private createUserQuotaData(quota: number, used: number): userQuotaData { const humanReadable = milisecondsToHumanReadable(this.lookback); return { @@ -91,7 +101,7 @@ export default class MessageCount implements IQuota { } } -function milisecondsToHumanReadable(totalMiliseconds: number): string { +export function milisecondsToHumanReadable(totalMiliseconds: number): string { const negative = totalMiliseconds < 0; if (negative) totalMiliseconds = -totalMiliseconds; diff --git a/src/quota/tokenCount.ts b/src/quota/tokenCount.ts new file mode 100644 index 0000000..f1ce0a3 --- /dev/null +++ b/src/quota/tokenCount.ts @@ -0,0 +1,124 @@ +import { User, GuildMember } from "discord.js"; + +import IQuota, { userQuotaData, userQuotaRecoveryData } from "../IQuota"; +import { apiRequest } from "../execution"; +import { database } from "../index"; +import { milisecondsToHumanReadable } from "./messageCount"; +import { Usage } from "@prisma/client"; + +/** + * Quota based on Tokens used and generated + */ +export default class tokenCount implements IQuota { + defaultQuota: number; + lookback: number; + considerInputTokensAsHalf: boolean; + + constructor( + defaultQuota: number = 512 * 25, + lookback: number = 1000 * 60 * 60 * 24, + considerInputTokensAsHalf: boolean = true, + ) { + this.defaultQuota = defaultQuota; + this.lookback = lookback; + this.considerInputTokensAsHalf = considerInputTokensAsHalf; + } + + private getUserQuota(id: string) { + return database.limits.findUnique({ + where: { user: BigInt(id) }, + }); + } + + async checkUser( + user: string | User | GuildMember, + request: apiRequest + ): Promise { + const userId: string = typeof user === "string" ? user : user.id; + + const userQuota = await this.getUserQuota(userId); + + const usedTokens = (await database.usage.aggregate({ + _sum: { + usageRequest: true, + usageResponse: true, + }, + where: { + user: BigInt(userId), + timestamp: { + gte: new Date(request.createdAt.getTime() - this.lookback), + } + } + }))._sum; + + if (!usedTokens.usageRequest || !usedTokens.usageResponse) throw new Error("Null from a database!! (tokenCount Quota)"); + + const usedUnits = (() => { + if (this.considerInputTokensAsHalf) + return usedTokens.usageResponse + usedTokens.usageRequest / 2; + return usedTokens.usageResponse + usedTokens.usageRequest; + })(); + + if (userQuota?.vip) return this.createUserQuotaData(Infinity, usedUnits); + if (userQuota?.limit) + return this.createUserQuotaData(userQuota.limit, usedUnits); + return this.createUserQuotaData(this.defaultQuota, usedUnits); + } + + findNthUsage(user: string, requestTimestamp: number, unitCount: number) { + if (this.considerInputTokensAsHalf) + throw("Not implemented"); + return database.$queryRaw>` + SELECT t1.*, ( + SELECT + SUM(usageResponse + usageRequest) 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 + `; + } + async getUserQuotaRecovery( + user: string | User | GuildMember, + request: apiRequest, + unitCount: number = 1 + ): Promise { + const userId = typeof user ==="string" ? user : user.id; + + const [userQuota, renameMebecause] = await Promise.all([ + this.checkUser(userId, request), + this.findNthUsage(userId, request.createdTimestamp, unitCount) + ]); + + console.log(renameMebecause); + return { + ...userQuota, + recoveryTimestamp: (renameMebecause.at(0)?.timestamp.valueOf() ?? Infinity) + this.lookback, + }; + } + + /** + * helper funtion to create userQuotaData + */ + private createUserQuotaData(quota: number, used: number): userQuotaData { + const humanReadable = milisecondsToHumanReadable(this.lookback); + return { + quota: quota, + used: used, + unitName: "tokens", + toString() { + return `${this.used} of ${this.quota} ${this.unitName} in ${humanReadable}`; + }, + }; + } +}