Quota: Refactor how Quotas are being handled
also renamed limits to quota I believe this new approach would allow me and bot hosters to add, implement or change the quota behavior more easily. Reimplemented the currently existing "Message count" limit to use the new IQuota, refactoring a code *a little*.
This commit is contained in:
parent
46bb5c867d
commit
06479472e5
5 changed files with 215 additions and 108 deletions
46
src/IQuota.ts
Normal file
46
src/IQuota.ts
Normal file
|
@ -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>;
|
||||||
|
}
|
|
@ -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 Command from "../command";
|
||||||
import { getUserLimit, getNthUseInLimitTimestamp } from "../execution";
|
import { config } from "../index";
|
||||||
|
|
||||||
export default class MyLimit extends Command implements Command {
|
export default class MyLimit extends Command implements Command {
|
||||||
name = "check-limit";
|
name = "check-limit";
|
||||||
|
@ -27,40 +27,23 @@ export default class MyLimit extends Command implements Command {
|
||||||
|
|
||||||
if (recoveryFor <= 0) recoveryFor = 1;
|
if (recoveryFor <= 0) recoveryFor = 1;
|
||||||
|
|
||||||
const userLimitPromise = getUserLimit(interaction.user, interaction.createdAt);
|
const userQuotaRecovery = await config.quota.getUserQuotaRecovery(interaction.user, interaction, recoveryFor);
|
||||||
const nthUseInLimitTimestampPromise = getNthUseInLimitTimestamp(
|
|
||||||
interaction.user,
|
|
||||||
interaction.createdAt,
|
|
||||||
recoveryFor,
|
|
||||||
);
|
|
||||||
|
|
||||||
const userLimit = await userLimitPromise;
|
const fields: APIEmbedField[] = [];
|
||||||
const nthUseInLimitTimestamp = await nthUseInLimitTimestampPromise;
|
|
||||||
|
|
||||||
if (userLimit === false || nthUseInLimitTimestamp === false) {
|
fields.push({ name: "Quota", inline: true, value: `${userQuotaRecovery.quota} ${userQuotaRecovery.unitName}`.trim() });
|
||||||
await interaction.reply({
|
fields.push({ name: "Usage", inline: true, value: `${userQuotaRecovery.used} ${userQuotaRecovery.unitName}`.trim() });
|
||||||
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;
|
|
||||||
}
|
|
||||||
|
|
||||||
await interaction.reply({
|
if (userQuotaRecovery.recoveryTimestamp !== undefined) fields.push({
|
||||||
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" : ""}`,
|
name: `Recovery for ${recoveryFor} message${recoveryFor>1 ? "s" : ""}`,
|
||||||
value: nthUseInLimitTimestamp === null ? "never" :
|
value: userQuotaRecovery.recoveryTimestamp === Infinity ? "never" :
|
||||||
// timestamp of the nth use in limit + 24 hours
|
`<t:${Math.ceil(userQuotaRecovery.recoveryTimestamp/1000)}:R>`
|
||||||
`<t:${Math.floor(new Date(nthUseInLimitTimestamp.getTime() + 1000 * 60 * 60 * 24 /*24 hours*/).getTime()/1000)}:R>`,
|
});
|
||||||
},
|
|
||||||
]
|
await interaction.reply({
|
||||||
|
embeds: [{
|
||||||
|
author: { name: interaction.user.username, icon_url: interaction.user.displayAvatarURL({ size: 128 }) },
|
||||||
|
fields: fields,
|
||||||
}],
|
}],
|
||||||
ephemeral: ephemeral,
|
ephemeral: ephemeral,
|
||||||
});
|
});
|
||||||
|
|
|
@ -4,6 +4,9 @@ import {
|
||||||
CreateChatCompletionRequest as ChatCompletionRequestData,
|
CreateChatCompletionRequest as ChatCompletionRequestData,
|
||||||
} from "openai";
|
} from "openai";
|
||||||
|
|
||||||
|
import IQuota from "./IQuota";
|
||||||
|
import MessageCount from "./quota/messageCount";
|
||||||
|
|
||||||
export interface IConfigRequired {
|
export interface IConfigRequired {
|
||||||
readonly calendarParams: Intl.DateTimeFormatOptions;
|
readonly calendarParams: Intl.DateTimeFormatOptions;
|
||||||
/** Tokens to authentiate with */
|
/** Tokens to authentiate with */
|
||||||
|
@ -24,11 +27,11 @@ export interface IConfigRequired {
|
||||||
/** Maximum total token usage for messages (counted locally) */
|
/** Maximum total token usage for messages (counted locally) */
|
||||||
readonly tokens: number;
|
readonly tokens: number;
|
||||||
};
|
};
|
||||||
/** Default limits for user inteacting with the bot */
|
/**
|
||||||
readonly userLimits: {
|
* Quota parameters to use when checking limits
|
||||||
/** How much requests can an user make if it's not overriden in a database entry */
|
* you can find some in `./quota`
|
||||||
readonly requests: number;
|
*/
|
||||||
};
|
readonly quota: IQuota;
|
||||||
}
|
}
|
||||||
|
|
||||||
export type IConfig = Partial<IConfigRequired>;
|
export type IConfig = Partial<IConfigRequired>;
|
||||||
|
@ -85,7 +88,5 @@ const defaultConfig: IConfigRequired = {
|
||||||
messages: envAsNumber("READ_LIMITS__MESSAGES") ?? 50,
|
messages: envAsNumber("READ_LIMITS__MESSAGES") ?? 50,
|
||||||
tokens: envAsNumber("READ_LIMITS__TOKENS") ?? 2048,
|
tokens: envAsNumber("READ_LIMITS__TOKENS") ?? 2048,
|
||||||
},
|
},
|
||||||
userLimits: {
|
quota: new MessageCount(envAsNumber("QUOTA__DEFAULT_QUOTA"), envAsNumber("QUOTA__LOOKBACK"))
|
||||||
requests: envAsNumber("USER_LIMITS__REQUESTS") ?? 25,
|
|
||||||
},
|
|
||||||
};
|
};
|
||||||
|
|
|
@ -8,7 +8,7 @@ import toOpenAIMessages from "./toOpenAIMessages";
|
||||||
import FunctionManager from "./funcitonManager";
|
import FunctionManager from "./funcitonManager";
|
||||||
|
|
||||||
type NonNullableInObject<T, V> = { [k in keyof T]: k extends V ? NonNullable<T[k]> : T[k] };
|
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">;
|
export type RequestMessage = apiRequest & NonNullableInObject<apiRequest, "channel" | "channelId">;
|
||||||
|
|
||||||
class ChannelsRunningValue extends Array<RequestMessage> {
|
class ChannelsRunningValue extends Array<RequestMessage> {
|
||||||
|
@ -60,69 +60,6 @@ export function getAuthor(request: apiRequest) {
|
||||||
return request.user;
|
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
|
* Replies to a request
|
||||||
* @param request the request to reply to
|
* @param request the request to reply to
|
||||||
|
@ -194,9 +131,9 @@ export async function queueRequest(request: apiRequest) {
|
||||||
return;
|
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) {
|
if (request instanceof DiscordApi.Message) {
|
||||||
request.react("🛑").catch(() => { /* NOTE: We send an informaton about limit reached in DM */ });
|
request.react("🛑").catch(() => { /* NOTE: We send an informaton about limit reached in DM */ });
|
||||||
if (!request.author.dmChannel) await request.author.createDM();
|
if (!request.author.dmChannel) await request.author.createDM();
|
||||||
|
@ -204,14 +141,13 @@ export async function queueRequest(request: apiRequest) {
|
||||||
embeds: [{
|
embeds: [{
|
||||||
color: 0xff0000,
|
color: 0xff0000,
|
||||||
description:
|
description:
|
||||||
"You've used up your message limit for today,\n" +
|
"You've used up your quota,\n" + userLimit.toString(),
|
||||||
`${userLimit.limit} requests in last 24 hours`,
|
|
||||||
}]
|
}]
|
||||||
}).catch(() => {/* FIXME: What should the bot do in this case to inform of limit reached?*/});
|
}).catch(() => {/* FIXME: What should the bot do in this case to inform of limit reached?*/});
|
||||||
}
|
}
|
||||||
else if (request.isRepliable()) {
|
else if (request.isRepliable()) {
|
||||||
request.reply({
|
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,
|
ephemeral: true,
|
||||||
}).catch(() => { /* Impossible to get there unless connection lost*/ });
|
}).catch(() => { /* Impossible to get there unless connection lost*/ });
|
||||||
}
|
}
|
||||||
|
|
141
src/quota/messageCount.ts
Normal file
141
src/quota/messageCount.ts
Normal file
|
@ -0,0 +1,141 @@
|
||||||
|
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;
|
||||||
|
|
||||||
|
constructor(
|
||||||
|
defaultQuota: number = 25,
|
||||||
|
lookback: number = 1000 * 60 * 60 * 24
|
||||||
|
) {
|
||||||
|
this.defaultQuota = defaultQuota;
|
||||||
|
this.lookback = lookback;
|
||||||
|
}
|
||||||
|
|
||||||
|
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,
|
||||||
|
};
|
||||||
|
}
|
||||||
|
|
||||||
|
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}`;
|
||||||
|
},
|
||||||
|
};
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
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;
|
||||||
|
}
|
Loading…
Reference in a new issue