Merge branch 'feature/quota'

src/IQuota.ts

@@ -0,0 +1,46 @@
+import { GuildMember, User } from "discord.js";
+
+import { apiRequest } from "./execution";
+
+export interface userQuotaData {
+  /** How much quota does the user have */
+  quota: number;
+  /** How much used quota does the user have */
+  used: number;
+  /** Name of the quota unit */
+  unitName?: string;
+  /** String representation of quota */
+  toString(): string;
+}
+
+export interface userQuotaRecoveryData extends userQuotaData {
+  /** Timestamp until asked QuotaRecovery will get fully recovered, Infinity if never */
+  recoveryTimestamp?: number;
+}
+/**
+ * Represents an quota engine that instance could use
+ */
+export default interface IQuota {
+  /**
+   * Gets user remining quota (or lack of it)
+   * @param request The request data
+   */
+  checkUser(
+    user: User | GuildMember | string,
+    request: apiRequest
+  ): Promise<userQuotaData>;
+  //checkUser(user: User | GuildMember | string, request: RequestMessage): userQuotaData;
+  /** Name of the Quota engine */
+  name?: string;
+  /**
+   * Checks recovery of quota units for a given user
+   * @param user the user to check
+   * @param request The request data
+   * @param unitCount for how many units to check for recovery
+   */
+  getUserQuotaRecovery(
+    user: User | GuildMember | string,
+    request: apiRequest,
+    unitCount?: number
+  ): Promise<userQuotaRecoveryData>;
+}

src/commands/check-limit.ts

@@ -1,7 +1,7 @@
-import { ApplicationCommandType, ChatInputCommandInteraction, ApplicationCommandOption, ApplicationCommandOptionType } from "discord.js";
+import { ApplicationCommandType, ChatInputCommandInteraction, ApplicationCommandOption, ApplicationCommandOptionType, APIEmbedField } from "discord.js";
 
 import Command from "../command";
-import { getUserLimit, getNthUseInLimitTimestamp } from "../execution";
+import { config } from "../index";
 
 export default class MyLimit extends Command implements Command {
   name = "check-limit";
@@ -27,40 +27,23 @@ export default class MyLimit extends Command implements Command {
 
     if (recoveryFor <= 0) recoveryFor = 1;
 
-    const userLimitPromise = getUserLimit(interaction.user, interaction.createdAt);
-    const nthUseInLimitTimestampPromise = getNthUseInLimitTimestamp(
-      interaction.user,
-      interaction.createdAt,
-      recoveryFor,
-    );
+    const userQuotaRecovery = await config.quota.getUserQuotaRecovery(interaction.user, interaction, recoveryFor);
 
-    const userLimit = await userLimitPromise;
-    const nthUseInLimitTimestamp = await nthUseInLimitTimestampPromise;
+    const fields: APIEmbedField[] = [];
 
-    if (userLimit === false || nthUseInLimitTimestamp === false) {
-      await interaction.reply({
-        embeds: [{
-          author: { name: interaction.user.username, icon_url: interaction.user.displayAvatarURL({ size: 128 }) },
-          description: "User is a VIP, so there is no limit",
-        }],
-        ephemeral: ephemeral,
-      });
-      return;
-    }
+    fields.push({ name: "Quota", inline: true, value: `${userQuotaRecovery.quota} ${userQuotaRecovery.unitName}`.trim() });
+    fields.push({ name: "Usage", inline: true, value: `${userQuotaRecovery.used} ${userQuotaRecovery.unitName}`.trim() });
+
+    if (userQuotaRecovery.recoveryTimestamp !== undefined) fields.push({
+      name: `Recovery for ${recoveryFor} ${userQuotaRecovery.unitName}`.trim(),
+      value: userQuotaRecovery.recoveryTimestamp === Infinity ? "never" :
+        `<t:${Math.ceil(userQuotaRecovery.recoveryTimestamp/1000)}:R>`
+    });
 
     await interaction.reply({
       embeds: [{
         author: { name: interaction.user.username, icon_url: interaction.user.displayAvatarURL({ size: 128 }) },
-        fields: [
-          { name: "Limit", inline: true, value: String(userLimit.limit) },
-          { name: "Usage", inline: true, value: String(userLimit.limit - userLimit.remaining) },
-          {
-            name: `Recovery for ${recoveryFor} message${recoveryFor>1 ? "s" : ""}`,
-            value: nthUseInLimitTimestamp === null ? "never" :
-                  // timestamp of the nth use in limit + 24 hours
-                  `<t:${Math.floor(new Date(nthUseInLimitTimestamp.getTime() + 1000 * 60 * 60 * 24 /*24 hours*/).getTime()/1000)}:R>`,
-          },
-        ]
+        fields: fields,
       }],
       ephemeral: ephemeral,
     });

src/configDefault.ts

@@ -4,6 +4,9 @@ import {
   CreateChatCompletionRequest as ChatCompletionRequestData,
 } from "openai";
 
+import IQuota from "./IQuota";
+import MessageCount from "./quota/messageCount";
+
 export interface IConfigRequired {
   readonly calendarParams: Intl.DateTimeFormatOptions;
   /** Tokens to authentiate with */
@@ -24,11 +27,11 @@ export interface IConfigRequired {
     /** Maximum total token usage for messages (counted locally) */
     readonly tokens: number;
   };
-  /** Default limits for user inteacting with the bot */
-  readonly userLimits: {
-    /** How much requests can an user make if it's not overriden in a database entry */
-    readonly requests: number;
-  };
+  /**
+   * Quota parameters to use when checking limits
+   * you can find some in `./quota`
+   */
+  readonly quota: IQuota;
 }
 
 export type IConfig = Partial<IConfigRequired>;
@@ -85,7 +88,5 @@ const defaultConfig: IConfigRequired = {
     messages: envAsNumber("READ_LIMITS__MESSAGES") ?? 50,
     tokens: envAsNumber("READ_LIMITS__TOKENS") ?? 2048,
   },
-  userLimits: {
-    requests: envAsNumber("USER_LIMITS__REQUESTS") ?? 25,
-  },
+  quota: new MessageCount(envAsNumber("QUOTA__DEFAULT_QUOTA"), envAsNumber("QUOTA__LOOKBACK"))
 };

src/execution.ts

@@ -8,7 +8,7 @@ import toOpenAIMessages from "./toOpenAIMessages";
 import FunctionManager from "./funcitonManager";
 
 type NonNullableInObject<T, V> = { [k in keyof T]: k extends V ? NonNullable<T[k]> : T[k] };
-type apiRequest = DiscordApi.Message | DiscordApi.RepliableInteraction;
+export type apiRequest = DiscordApi.Message | DiscordApi.RepliableInteraction;
 export type RequestMessage = apiRequest & NonNullableInObject<apiRequest, "channel" | "channelId">;
 
 class ChannelsRunningValue extends Array<RequestMessage> {
@@ -60,69 +60,6 @@ export function getAuthor(request: apiRequest) {
   return request.user;
 }
 
-/**
- * gets user remaining limit (or lack of it)
- * @param user the user to check
- * @param requestTimestamp the timestamp of the user request
- * @returns object containing the limit and remaining usage or `false` if there is no limit
- */
-export async function getUserLimit(user: string | { id: string }, requestTimestamp: Date) {
-  const userId: string = typeof user === "string" ? user : user.id;
-
-  const userLimits = await database.limits.findUnique({
-    where: { user: BigInt(userId) }
-  });
-
-  if (userLimits?.vip) return false;
-
-  const usedLimit = (await database.usage.count({
-    select: { _all: true },
-    where: {
-      user: BigInt(userId),
-      timestamp: {
-        gte: new Date(requestTimestamp.getTime() - 1000 * 60 * 60 * 24 /* 24 hours */)
-      }
-    },
-  }))._all;
-
-  if (!userLimits || !userLimits.limit) return {limit: config.userLimits.requests, remaining: config.userLimits.requests - usedLimit};
-
-  return {limit: userLimits.limit, remaining: userLimits.limit - usedLimit};
-}
-
-/**
- * gets the timestamp of nth use inside time limit
- * @param user the user or id to check
- * @param requestTimestamp the timestamp of the request (message/interaction createdAt)
- * @param nth which timestamp in time limit to get (orderedd from oldest to newest)
- * @returns `false` if user is vip
- * @returns `null` if there is no request
- * @returns `Date` timestamp of the nth request
- */
-export async function getNthUseInLimitTimestamp(user: string | { id: string }, requestTimestamp: Date, nth = 1) {
-  const userId: string = typeof user === "string" ? user : user.id;
-
-  const userLimits = await database.limits.findUnique({
-    where: { user: BigInt(userId)}
-  });
-
-  if (userLimits?.vip) return false;
-
-  const nthUseInLimit = await database.usage.findFirst({
-    where: {
-      user: BigInt(userId),
-      timestamp: {
-        gte: new Date(requestTimestamp.getTime() - 1000 * 60 * 60 * 24 /* 24 hours */)
-      }
-    },
-    orderBy: { timestamp: "asc" },
-    skip: nth - 1,
-  });
-
-  if (!nthUseInLimit) return null;
-  return nthUseInLimit.timestamp;
-}
-
 /**
  * Replies to a request
  * @param request the request to reply to
@@ -194,9 +131,9 @@ export async function queueRequest(request: apiRequest) {
     return;
   }
 
-  const userLimit = await getUserLimit(getAuthor(request), request.createdAt);
+  const userLimit = await config.quota.checkUser(getAuthor(request), request);
 
-  if (userLimit !== false && userLimit.remaining <= 0) {
+  if (userLimit.used >= userLimit.quota) {
     if (request instanceof DiscordApi.Message) {
       request.react("🛑").catch(() => { /* NOTE: We send an informaton about limit reached in DM */ });
       if (!request.author.dmChannel) await request.author.createDM();
@@ -204,14 +141,13 @@ export async function queueRequest(request: apiRequest) {
         embeds: [{
           color: 0xff0000,
           description:
-            "You've used up your message limit for today,\n" +
-            `${userLimit.limit} requests in last 24 hours`,
+            "You've used up your quota,\n" + userLimit.toString(),
         }]
       }).catch(() => {/* FIXME: What should the bot do in this case to inform of limit reached?*/});
     }
     else if (request.isRepliable()) {
       request.reply({
-        content: `You've used up your message limit for today, ${userLimit.limit} requests in last 24 hours`,
+        content: "You've used up your quota\n" + userLimit.toString(),
         ephemeral: true,
       }).catch(() => { /* Impossible to get there unless connection lost*/ });
     }

src/quota/messageCount.ts

@@ -0,0 +1,151 @@
+import { GuildMember, User } from "discord.js";
+
+import IQuota, { userQuotaData, userQuotaRecoveryData } from "../IQuota";
+import { apiRequest } from "../execution";
+import { database } from "../index";
+
+/**
+ * Quota based on Request message count and function calls
+ */
+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
+  ) {
+    this.defaultQuota = defaultQuota;
+    this.lookback = lookback;
+  }
+
+  /**
+   * Retrives the quota from the database
+   */
+  private getUserQuota(id: string) {
+    return database.limits.findUnique({
+      where: { user: BigInt(id) },
+    });
+  }
+
+  async checkUser(
+    user: User | GuildMember | string,
+    request: apiRequest
+  ): Promise<userQuotaData> {
+    const userId: string = typeof user === "string" ? user : user.id;
+
+    const userQuota = await this.getUserQuota(userId);
+
+    const usedLimit = (
+      await database.usage.count({
+        select: { _all: true },
+        where: {
+          user: BigInt(userId),
+          timestamp: {
+            gte: new Date(request.createdAt.getTime() - this.lookback),
+          },
+        },
+      })
+    )._all;
+
+    if (userQuota?.vip) return this.createUserQuotaData(Infinity, usedLimit);
+    if (userQuota?.limit)
+      return this.createUserQuotaData(userQuota.limit, usedLimit);
+    return this.createUserQuotaData(this.defaultQuota, usedLimit);
+  }
+
+  async getUserQuotaRecovery(
+    user: string | User | GuildMember,
+    request: apiRequest,
+    unitCount: number | undefined = 1
+  ): Promise<userQuotaRecoveryData> {
+    const userId = typeof user === "string" ? user : user.id;
+
+    const [userQuota, nthUseInLimit] = await Promise.all([
+      this.checkUser(userId, request),
+      database.usage.findFirst({
+        where: {
+          user: BigInt(userId),
+          timestamp: {
+            gte: new Date(request.createdTimestamp - this.lookback),
+          },
+        },
+        orderBy: { timestamp: "asc" },
+        skip: unitCount - 1,
+      }),
+    ]);
+
+    return {
+      ...userQuota,
+      recoveryTimestamp: nthUseInLimit ? this.lookback + nthUseInLimit?.timestamp.valueOf() : Infinity,
+    };
+  }
+
+  /**
+   * helper funtion to create userQuotaData
+   */
+  private createUserQuotaData(quota: number, used: number): userQuotaData {
+    const humanReadable = milisecondsToHumanReadable(this.lookback);
+    return {
+      quota: quota,
+      used: used,
+      unitName: "messages and model function calls",
+      toString() {
+        return `${this.used} of ${this.quota} ${this.unitName} in ${humanReadable}`;
+      },
+    };
+  }
+}
+
+export function milisecondsToHumanReadable(totalMiliseconds: number): string {
+  const negative = totalMiliseconds < 0;
+  if (negative) totalMiliseconds = -totalMiliseconds;
+
+  const totalSeconds = totalMiliseconds / 1000;
+  const totalMinutes = totalSeconds / 60;
+  const totalHours = totalMinutes / 60;
+  const totalDays = totalHours / 24;
+  const totalYears = totalDays / 365.2425;
+
+  //const hours = Math.floor(totalHours % 24);
+  const minutes = Math.floor(totalMinutes % 60);
+  const seconds = Math.floor(totalSeconds % 60);
+  const miliseconds = Math.floor(totalMiliseconds % 1000);
+
+  if (totalYears > 1.01) return `${Math.ceil(totalYears * 100) / 100} years`;
+  if (totalDays > 16) return `${Math.ceil(totalDays * 100) / 100} days`;
+  if (totalDays > 1 && totalDays % 1 === 0) return `${totalDays} days`;
+  if (totalDays > 4 && totalHours % 1 === 0)
+    return `${Math.floor(totalDays)} days ${
+      Math.ceil((totalHours % 24) * 100) / 100
+    } hours`;
+
+  const textHour = totalHours === 0 ? "" : totalHours === 1 ? "1 hour" : `${totalHours} hours`;
+  const textMinutes =
+    minutes === 0 ? "" : minutes === 1 ? "1 minute" : `${minutes} minutes`;
+  const textSeconds =
+    seconds === 0 ? "" : seconds === 1 ? "1 second" : `${seconds} seconds`;
+  const textMiliseconds =
+    miliseconds === 0
+      ? ""
+      : miliseconds === 1
+      ? "1 milisecond"
+      : `${Math.ceil(miliseconds)} miliseconds`;
+  const textMilisecondsPrecise =
+    miliseconds === 0
+      ? ""
+      : miliseconds === 1
+      ? "1 milisecond"
+      : `${Math.ceil(miliseconds * 1e6) / 1e6} miliseconds`;
+
+  if (totalHours >= 1) return [textHour, textMinutes].join(" ").trim();
+  if (totalMinutes >= 1) return [textMinutes, textSeconds].join(" ").trim();
+  if (totalSeconds >= 10)
+    return [textSeconds, textMiliseconds].join(" ").trim();
+  if (totalMiliseconds >= 10) return [textMiliseconds].join(" ").trim();
+  return textMilisecondsPrecise;
+}

src/quota/tokenCount.ts

@@ -0,0 +1,148 @@
+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}`;
+      },
+    };
+  }
+}