diff --git a/src/IQuota.ts b/src/IQuota.ts new file mode 100644 index 0000000..2dbf23c --- /dev/null +++ b/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; + //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; +} diff --git a/src/commands/check-limit.ts b/src/commands/check-limit.ts index 870874f..c9e667b 100644 --- a/src/commands/check-limit.ts +++ b/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" : + `` + }); 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 - ``, - }, - ] + fields: fields, }], ephemeral: ephemeral, }); diff --git a/src/configDefault.ts b/src/configDefault.ts index a4a6bf3..608806b 100644 --- a/src/configDefault.ts +++ b/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; @@ -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")) }; diff --git a/src/execution.ts b/src/execution.ts index a755ec6..aa99e73 100644 --- a/src/execution.ts +++ b/src/execution.ts @@ -8,7 +8,7 @@ import toOpenAIMessages from "./toOpenAIMessages"; import FunctionManager from "./funcitonManager"; type NonNullableInObject = { [k in keyof T]: k extends V ? NonNullable : T[k] }; -type apiRequest = DiscordApi.Message | DiscordApi.RepliableInteraction; +export type apiRequest = DiscordApi.Message | DiscordApi.RepliableInteraction; export type RequestMessage = apiRequest & NonNullableInObject; class ChannelsRunningValue extends Array { @@ -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*/ }); } diff --git a/src/quota/messageCount.ts b/src/quota/messageCount.ts new file mode 100644 index 0000000..7b97755 --- /dev/null +++ b/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 { + 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 { + 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; +} diff --git a/src/quota/tokenCount.ts b/src/quota/tokenCount.ts new file mode 100644 index 0000000..263ba13 --- /dev/null +++ b/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 { + 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>` + 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>` + 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) + ]); + + 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}`; + }, + }; + } +}