GPTcord/src/quota/tokenCount.ts

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);
  }

  /**
   * Helper function to use proper query based on considerInputTokensAsHalf setting
   * @param user The user to find nth usage
   * @param requestTimestamp the timestamp of the request
   * @param unitCount the unit count to check
   * @returns promise of giving out the record
   */
  findNthUsage(user: string, requestTimestamp: number, unitCount: number) {
    if (this.considerInputTokensAsHalf)
      return database.$queryRaw<Array<Usage & {usage: number}>>`
        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<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)
    ]);

    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}`;
      },
    };
  }
}