/check-limit: fix wrong unit name in field name

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" :
         `<t:${Math.ceil(userQuotaRecovery.recoveryTimestamp/1000)}:R>`
     });

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;
 

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<userQuotaData> {
+    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<Array<Usage & {usage: bigint}>>`
+      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<userQuotaRecoveryData> {
+    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}`;
+      },
+    };
+  }
+}