diff --git a/.github/PULL_REQUEST_TEMPLATE.md b/.github/PULL_REQUEST_TEMPLATE.md index fc790f0..a38516c 100644 --- a/.github/PULL_REQUEST_TEMPLATE.md +++ b/.github/PULL_REQUEST_TEMPLATE.md @@ -50,5 +50,5 @@ Issue: # ## Discord User (optional) - - \ No newline at end of file + + \ No newline at end of file diff --git a/packages/discord/README.md b/packages/discord/README.md index 034388e..0e98dbc 100644 --- a/packages/discord/README.md +++ b/packages/discord/README.md @@ -64,9 +64,9 @@ This is done through environment variables. Refer to the [`.env.example`](./.env > The Quota System is opt-in. Specifying a database connection string in your environment variables will enable it automatically. See the [`.env.example`](./.env.example) file for more information. -On top of whitelisting and blacklisting, you can optionally enable the bot's quota system. This system allows you to specify a maximum amount of [tokens](https://openai.com/pricing) that can be used for each allowed user. It uses key-value storage powered by [KeyV](https://github.com/jaredwray/keyv), which is why multiple databases are supported (MongoDB, PostgreSQL and MySQL). +On top of whitelisting and blacklisting, you can optionally enable the bot's quota system. This system allows you to specify a maximum amount of [tokens](https://openai.com/pricing) that can be used for each allowed user or role. It uses key-value storage powered by [KeyV](https://github.com/jaredwray/keyv), which is why multiple databases are supported (MongoDB, PostgreSQL and MySQL). -> If you haven't explicitly set a quota for a user, the bot will use the default quota of 5000 tokens (about $0.01). This can be changed by setting the `DEFAULT_QUOTA` environment variable. +> If you haven't explicitly set a quota for a user or one of the user's roles, the bot will use the default quota of 5000 tokens (about $0.01). This can be changed by setting the `DEFAULT_QUOTA` environment variable. Important things to note about the quota system: @@ -75,6 +75,7 @@ Important things to note about the quota system: - The tokens are calculated by `gpt-turbo` using third-party libraries. The calculated tokens may *slightly* differ from OpenAI's actual token count. This shouldn't be too much of an issue for small quotas (e.g. 5000 tokens), but it may be an issue for larger quotas (e.g. 1000000 tokens). - The bot will set the `max_tokens` accordingly to prevent the user from going over their quota. - For example: A user's usage is at 5/10 tokens and their current conversation size is 2 tokens, this means whatever prompt is next will incur a cost of at least `5 + 2 = 7` tokens. The bot will set the `max_tokens` parameter to `10 - 7 = 3` tokens to prevent the user from going over their quota. Again, depending on how OpenAI actually respects this parameter, there may be a slight difference in the actual token count. +- User quotas have priority over role quotas. ## The `MESSAGE_CONTENT` intent diff --git a/packages/discord/src/GPTTurboClient.ts b/packages/discord/src/GPTTurboClient.ts index 1d07655..d646aec 100644 --- a/packages/discord/src/GPTTurboClient.ts +++ b/packages/discord/src/GPTTurboClient.ts @@ -20,9 +20,10 @@ import SlashCommandHandler from "./interaction-handlers/SlashCommandHandler.js"; import AdminUsageResetMenuHandler from "./interaction-handlers/admin/usage-reset/AdminUsageResetMenuHandler.js"; import AdminUsageResetAllHandler from "./interaction-handlers/admin/usage-reset/AdminUsageResetAllHandler.js"; import AdminUsageResetUserHandler from "./interaction-handlers/admin/usage-reset/AdminUsageResetUserHandler.js"; -import AdminQuotaUserMenuHandler from "./interaction-handlers/admin/quota/AdminQuotaUserMenuHandler.js"; +import AdminQuotaMenuHandler from "./interaction-handlers/admin/quota/AdminQuotaMenuHandler.js"; import AdminQuotaUserHandler from "./interaction-handlers/admin/quota/AdminQuotaUserHandler.js"; -import UserIdModalHandler from "./interaction-handlers/UserIdModalHandler.js"; +import SnowflakeModalHandler from "./interaction-handlers/SnowflakeModalHandler.js"; +import AdminQuotaRoleHandler from "./interaction-handlers/admin/quota/AdminQuotaRoleHandler.js"; export default class GPTTurboClient< Ready extends boolean = boolean @@ -30,7 +31,7 @@ export default class GPTTurboClient< public id: string; public commandManager = new CommandManager(); public eventManager: EventManager = new EventManager(this); - public quotaManager = new QuotaManager(); + public quotaManager: QuotaManager = new QuotaManager(this); public conversationManager = new ConversationManager(this.quotaManager); private cooldowns = new Collection>(); @@ -50,16 +51,27 @@ export default class GPTTurboClient< new AdminUsageResetMenuHandler(), new AdminUsageResetAllHandler(), new AdminUsageResetUserHandler(), - new UserIdModalHandler( + new SnowflakeModalHandler( AdminUsageResetMenuHandler.USER_ID_MODAL_ID, - AdminUsageResetUserHandler.ID + AdminUsageResetUserHandler.ID, + "Enter User ID", + "User ID" ), - new AdminQuotaUserMenuHandler(), + new AdminQuotaMenuHandler(), new AdminQuotaUserHandler(), - new UserIdModalHandler( - AdminQuotaUserMenuHandler.USER_ID_MODAL_ID, - AdminQuotaUserHandler.ID + new SnowflakeModalHandler( + AdminQuotaMenuHandler.USER_ID_MODAL_ID, + AdminQuotaUserHandler.ID, + "Enter User ID", + "User ID" + ), + new AdminQuotaRoleHandler(), + new SnowflakeModalHandler( + AdminQuotaMenuHandler.ROLE_ID_MODAL_ID, + AdminQuotaRoleHandler.ID, + "Enter Role ID", + "Role ID" ) ); @@ -107,7 +119,9 @@ export default class GPTTurboClient< async login(token = DISCORD_TOKEN) { await this.init(); - return super.login(token); + const res = await super.login(token); + await this.postInit(); + return res; } private async init() { @@ -117,4 +131,8 @@ export default class GPTTurboClient< this.quotaManager.init(), ]); } + + private async postInit() { + await Promise.all([this.quotaManager.postInit()]); + } } diff --git a/packages/discord/src/commands/admin.ts b/packages/discord/src/commands/admin.ts index 968d649..77ae7a0 100644 --- a/packages/discord/src/commands/admin.ts +++ b/packages/discord/src/commands/admin.ts @@ -12,7 +12,7 @@ import { ADMIN_MENU_ID } from "../config/constants.js"; import AdminUsageResetMenuHandler from "../interaction-handlers/admin/usage-reset/AdminUsageResetMenuHandler.js"; import setupInteractionCleanup from "../utils/setupInteractionCleanup.js"; import reply from "../utils/reply.js"; -import AdminQuotaUserMenuHandler from "../interaction-handlers/admin/quota/AdminQuotaUserMenuHandler.js"; +import AdminQuotaUserMenuHandler from "../interaction-handlers/admin/quota/AdminQuotaMenuHandler.js"; const adminCommand: DiscordSlashCommand = { builder: new SlashCommandBuilder() diff --git a/packages/discord/src/commands/usage.ts b/packages/discord/src/commands/usage.ts index 0bfd20a..fe0b113 100644 --- a/packages/discord/src/commands/usage.ts +++ b/packages/discord/src/commands/usage.ts @@ -16,16 +16,17 @@ const usageCommand: DiscordSlashCommand = { if (!interaction.isRepliable()) return; const { quotaManager } = interaction.client; - const [quota, usage] = await Promise.all([ + const [quota, usage, [quotaRole]] = await Promise.all([ quotaManager.getQuota(interaction.user.id), quotaManager.getUsage(interaction.user.id), + quotaManager.getUserQuotaRole(interaction.user.id), ]); - const left = quota - usage; + const remaining = quota - usage; const percent = usage / quota; const quotaTokens = tokenFormat.format(quota); const usageTokens = tokenFormat.format(Math.min(usage, quota)); - const leftTokens = tokenFormat.format(left); + const remainingTokens = tokenFormat.format(remaining); const percentString = percentFormat.format(Math.min(percent, 1)); const isExceeded = usage >= quota; @@ -40,10 +41,29 @@ const usageCommand: DiscordSlashCommand = { embeds: [ { title: `Usage (${percentString})`, - description: - left <= 0 - ? `You have used ${usageTokens} tokens out of your ${quotaTokens} quota` - : `You have used ${usageTokens} tokens out of your ${quotaTokens} quota (${leftTokens} left)`, + fields: [ + { + name: "Usage", + value: `${usageTokens}`, + inline: true, + }, + { + name: "Quota", + value: `${quotaTokens}`, + inline: true, + }, + { + name: "Remaining", + value: `${remainingTokens}`, + inline: true, + }, + ], + footer: + quotaRole !== null + ? { + text: `❗ Quota given for having the "${quotaRole.name}" role`, + } + : undefined, color, }, ], diff --git a/packages/discord/src/components/RoleSelect.ts b/packages/discord/src/components/RoleSelect.ts new file mode 100644 index 0000000..1687df9 --- /dev/null +++ b/packages/discord/src/components/RoleSelect.ts @@ -0,0 +1,14 @@ +import { RoleSelectMenuBuilder } from "discord.js"; +import getHandlerId from "../utils/getHandlerId.js"; + +export default class RoleSelect extends RoleSelectMenuBuilder { + public static readonly GENERIC_ID = getHandlerId(RoleSelect.name); + + constructor(id: string) { + super(); + this.setCustomId(id) + .setPlaceholder("Select a role") + .setMinValues(1) + .setMaxValues(1); + } +} diff --git a/packages/discord/src/interaction-handlers/UserIdModalHandler.ts b/packages/discord/src/interaction-handlers/SnowflakeModalHandler.ts similarity index 75% rename from packages/discord/src/interaction-handlers/UserIdModalHandler.ts rename to packages/discord/src/interaction-handlers/SnowflakeModalHandler.ts index d3473a2..13a2c19 100644 --- a/packages/discord/src/interaction-handlers/UserIdModalHandler.ts +++ b/packages/discord/src/interaction-handlers/SnowflakeModalHandler.ts @@ -9,20 +9,23 @@ import InteractionHandler from "./InteractionHandler.js"; import isBotAdmin from "../utils/isBotAdmin.js"; import getHandlerId from "../utils/getHandlerId.js"; -export default class UserIdModalHandler extends InteractionHandler { - public static readonly GENERIC_ID = getHandlerId(UserIdModalHandler.name); - public static readonly INPUT_ID = `${UserIdModalHandler.GENERIC_ID}_input`; +export default class SnowflakeModalHandler extends InteractionHandler { + public static readonly GENERIC_ID = getHandlerId( + SnowflakeModalHandler.name + ); + public static readonly INPUT_ID = `${SnowflakeModalHandler.GENERIC_ID}_input`; constructor( private id: string, private forwardId: string, - private title = "Enter User ID" + private title = "Enter ID", + private label = "ID" ) { super(); } public get name(): string { - return UserIdModalHandler.name; + return SnowflakeModalHandler.name; } protected canHandle(interaction: Interaction): Awaitable { @@ -43,8 +46,8 @@ export default class UserIdModalHandler extends InteractionHandler { .addComponents( this.createModalActionRow().addComponents( new TextInputBuilder() - .setCustomId(UserIdModalHandler.INPUT_ID) - .setLabel("User ID") + .setCustomId(SnowflakeModalHandler.INPUT_ID) + .setLabel(this.label) .setPlaceholder("123456789012345678") .setMinLength(18) .setMaxLength(18) diff --git a/packages/discord/src/interaction-handlers/admin/quota/AdminQuotaUserMenuHandler.ts b/packages/discord/src/interaction-handlers/admin/quota/AdminQuotaMenuHandler.ts similarity index 58% rename from packages/discord/src/interaction-handlers/admin/quota/AdminQuotaUserMenuHandler.ts rename to packages/discord/src/interaction-handlers/admin/quota/AdminQuotaMenuHandler.ts index 2710fc1..3cbb007 100644 --- a/packages/discord/src/interaction-handlers/admin/quota/AdminQuotaUserMenuHandler.ts +++ b/packages/discord/src/interaction-handlers/admin/quota/AdminQuotaMenuHandler.ts @@ -2,6 +2,8 @@ import { Interaction, Awaitable, StringSelectMenuInteraction, + ButtonBuilder, + ButtonStyle, } from "discord.js"; import InteractionHandler from "../../InteractionHandler.js"; import isBotAdmin from "../../../utils/isBotAdmin.js"; @@ -12,13 +14,16 @@ import reply from "../../../utils/reply.js"; import AdminQuotaUserHandler from "./AdminQuotaUserHandler.js"; import UserIdButton from "../../../components/UserIdButton.js"; import UserSelect from "../../../components/UserSelect.js"; +import RoleSelect from "../../../components/RoleSelect.js"; +import AdminQuotaRoleHandler from "./AdminQuotaRoleHandler.js"; -export default class AdminQuotaUserMenuHandler extends InteractionHandler { - public static readonly ID = getHandlerId(AdminQuotaUserMenuHandler.name); - public static readonly USER_ID_MODAL_ID = `${AdminQuotaUserMenuHandler.ID}_user-id-modal`; +export default class AdminQuotaMenuHandler extends InteractionHandler { + public static readonly ID = getHandlerId(AdminQuotaMenuHandler.name); + public static readonly USER_ID_MODAL_ID = `${AdminQuotaMenuHandler.ID}_user-id-modal`; + public static readonly ROLE_ID_MODAL_ID = `${AdminQuotaMenuHandler.ID}_role-id-modal`; public get name(): string { - return AdminQuotaUserMenuHandler.name; + return AdminQuotaMenuHandler.name; } protected canHandle(interaction: Interaction): Awaitable { @@ -26,7 +31,7 @@ export default class AdminQuotaUserMenuHandler extends InteractionHandler { isBotAdmin(interaction.user.id) && interaction.isStringSelectMenu() && interaction.customId === ADMIN_MENU_ID && - interaction.values[0] === AdminQuotaUserMenuHandler.ID + interaction.values[0] === AdminQuotaMenuHandler.ID ); } @@ -41,7 +46,15 @@ export default class AdminQuotaUserMenuHandler extends InteractionHandler { new UserSelect(AdminQuotaUserHandler.ID) ), this.createMessageActionRow().addComponents( - new UserIdButton(AdminQuotaUserMenuHandler.USER_ID_MODAL_ID) + new RoleSelect(AdminQuotaRoleHandler.ID) + ), + this.createMessageActionRow().addComponents( + new UserIdButton(AdminQuotaMenuHandler.USER_ID_MODAL_ID), + new ButtonBuilder() + .setCustomId(AdminQuotaMenuHandler.ROLE_ID_MODAL_ID) + .setLabel("Enter Role ID") + .setEmoji("🆔") + .setStyle(ButtonStyle.Secondary) ), ], }); diff --git a/packages/discord/src/interaction-handlers/admin/quota/AdminQuotaRoleHandler.ts b/packages/discord/src/interaction-handlers/admin/quota/AdminQuotaRoleHandler.ts new file mode 100644 index 0000000..4833339 --- /dev/null +++ b/packages/discord/src/interaction-handlers/admin/quota/AdminQuotaRoleHandler.ts @@ -0,0 +1,113 @@ +import { + Interaction, + Colors, + italic, + ButtonBuilder, + ButtonStyle, + InteractionResponse, + Message, +} from "discord.js"; +import getHandlerId from "../../../utils/getHandlerId.js"; +import reply from "../../../utils/reply.js"; +import AdminQuotaSnowflakeHandler from "./AdminQuotaSnowflakeHandler.js"; + +export default class AdminQuotaRoleHandler extends AdminQuotaSnowflakeHandler { + public static readonly ID = getHandlerId(AdminQuotaRoleHandler.name); + + constructor() { + super(AdminQuotaRoleHandler.ID); + } + + public get name(): string { + return AdminQuotaRoleHandler.name; + } + + protected async getCurrentQuota( + interaction: Interaction, + roleId: string + ): Promise { + const { quotaManager } = interaction.client; + return ( + (await quotaManager.getRoleQuota(roleId)) ?? + (await quotaManager.getDefaultQuota()) + ); + } + + protected async getInitialReply( + interaction: Interaction, + roleId: string + ): Promise | InteractionResponse | null> { + const { quotaManager } = interaction.client; + + const quota = await this.getCurrentQuota(interaction, roleId); + const hasQuota = await quotaManager.hasRoleQuota(roleId); + + const quotaFormat = Intl.NumberFormat("en-US", { + style: "decimal", + maximumFractionDigits: 0, + }).format(quota); + const quotaType = (() => { + if (hasQuota) { + return ""; + } + return `(${italic("default")})`; + })(); + + const row = this.createMessageActionRow().addComponents( + new ButtonBuilder() + .setCustomId(AdminQuotaSnowflakeHandler.BUTTON_SET_ID) + .setLabel("Set Quota") + .setStyle(ButtonStyle.Primary) + .setEmoji("🔧") + ); + + if (hasQuota) { + row.addComponents( + new ButtonBuilder() + .setCustomId(AdminQuotaSnowflakeHandler.BUTTON_RESET_ID) + .setLabel("Reset Quota") + .setStyle(ButtonStyle.Primary) + .setEmoji("🔄") + ); + } + + return reply(interaction, { + ephemeral: true, + embeds: [ + { + title: "Role Quota Information", + fields: [ + { + name: "Current Quota", + value: `${quotaFormat} ${quotaType}`, + inline: true, + }, + ], + color: Colors.Blue, + }, + ], + components: [row], + }); + } + + protected async deleteQuota( + interaction: Interaction, + roleId: string + ): Promise { + const { quotaManager } = interaction.client; + await quotaManager.deleteRoleQuota(roleId); + } + + protected getModalTitle(): string { + return "Set Role Quota"; + } + + protected async setQuota( + interaction: Interaction, + roleId: string, + quota: number + ): Promise { + const { quotaManager } = interaction.client; + await quotaManager.setRoleQuota(roleId, quota); + } +} diff --git a/packages/discord/src/interaction-handlers/admin/quota/AdminQuotaSnowflakeHandler.ts b/packages/discord/src/interaction-handlers/admin/quota/AdminQuotaSnowflakeHandler.ts new file mode 100644 index 0000000..baff61d --- /dev/null +++ b/packages/discord/src/interaction-handlers/admin/quota/AdminQuotaSnowflakeHandler.ts @@ -0,0 +1,216 @@ +import { + Interaction, + Awaitable, + ModalSubmitInteraction, + Colors, + DiscordjsError, + DiscordjsErrorCodes, + RepliableInteraction, + ModalBuilder, + ButtonInteraction, + TextInputStyle, + AnySelectMenuInteraction, + Message, + InteractionResponse, +} from "discord.js"; +import InteractionHandler from "../../InteractionHandler.js"; +import isBotAdmin from "../../../utils/isBotAdmin.js"; +import getHandlerId from "../../../utils/getHandlerId.js"; +import reply from "../../../utils/reply.js"; +import setupInteractionCleanup from "../../../utils/setupInteractionCleanup.js"; +import { TextInputBuilder } from "@discordjs/builders"; +import BotException from "../../../exceptions/BotException.js"; +import { DEFAULT_INTERACTION_WAIT } from "../../../config/constants.js"; +import SnowflakeModalHandler from "../../SnowflakeModalHandler.js"; + +export default abstract class AdminQuotaSnowflakeHandler extends InteractionHandler { + public static readonly GENERIC_ID = getHandlerId( + AdminQuotaSnowflakeHandler.name + ); + protected static readonly BUTTON_SET_ID = `${AdminQuotaSnowflakeHandler.GENERIC_ID}_button-set`; + protected static readonly BUTTON_RESET_ID = `${AdminQuotaSnowflakeHandler.GENERIC_ID}_button-reset`; + protected static readonly MODAL_SET_ID = `${AdminQuotaSnowflakeHandler.GENERIC_ID}_modal-set`; + protected static readonly QUOTA_INPUT_ID = `${AdminQuotaSnowflakeHandler.GENERIC_ID}_quota-input`; + + constructor(private id: string) { + super(); + } + + protected canHandle(interaction: Interaction): Awaitable { + return ( + isBotAdmin(interaction.user.id) && + (interaction.isAnySelectMenu() || interaction.isModalSubmit()) && + interaction.customId === this.id + ); + } + + protected async handle( + interaction: AnySelectMenuInteraction | ModalSubmitInteraction + ): Promise { + const snowflake = interaction.isModalSubmit() + ? interaction.fields.getTextInputValue( + SnowflakeModalHandler.INPUT_ID + ) + : interaction.values[0]; + + const response = await this.getInitialReply(interaction, snowflake); + if (!response) throw new Error("Failed to send message"); + + try { + const buttonInteraction = await response.awaitMessageComponent({ + filter: (i) => + i.user.id === interaction.user.id && + [ + AdminQuotaSnowflakeHandler.BUTTON_SET_ID, + AdminQuotaSnowflakeHandler.BUTTON_RESET_ID, + ].includes(i.customId), + time: DEFAULT_INTERACTION_WAIT, + }); + await response.delete(); + + if (!buttonInteraction.isButton()) + throw new Error("Expected button interaction"); + + await this.handleButtonInteraction(buttonInteraction, snowflake); + } catch (e) { + this.handleError(interaction, e); + } + } + + private async handleButtonInteraction( + interaction: ButtonInteraction, + snowflake: string + ) { + switch (interaction.customId) { + case AdminQuotaSnowflakeHandler.BUTTON_SET_ID: + await this.showQuotaModal(interaction, snowflake); + break; + case AdminQuotaSnowflakeHandler.BUTTON_RESET_ID: + await this.deleteQuota(interaction, snowflake); + const quota = await this.getCurrentQuota( + interaction, + snowflake + ); + await reply(interaction, { + ephemeral: true, + embeds: [ + { + title: "Success", + description: `Quota was reset to ${quota}`, + color: Colors.Green, + }, + ], + }); + break; + default: + throw new Error( + `Unknown button custom id: ${interaction.customId}` + ); + } + } + + private async showQuotaModal( + interaction: ButtonInteraction, + snowflake: string + ) { + const quota = await this.getCurrentQuota(interaction, snowflake); + + const modal = new ModalBuilder() + .setCustomId(AdminQuotaSnowflakeHandler.MODAL_SET_ID) + .setTitle(this.getModalTitle()) + .addComponents( + this.createModalActionRow().addComponents( + new TextInputBuilder() + .setCustomId(AdminQuotaSnowflakeHandler.QUOTA_INPUT_ID) + .setLabel("Quota") + .setPlaceholder(quota.toString()) + .setRequired(true) + .setMinLength(1) + .setMaxLength(10) + .setStyle(TextInputStyle.Short) + .setValue(quota.toString()) + ) + ); + + await interaction.showModal(modal); + + const modalInteraction = await interaction.awaitModalSubmit({ + filter: (i) => + i.user.id === interaction.user.id && + i.isModalSubmit() && + i.customId === AdminQuotaSnowflakeHandler.MODAL_SET_ID, + time: DEFAULT_INTERACTION_WAIT, + }); + await modalInteraction.deferUpdate(); + + const quotaInput = modalInteraction.fields.getTextInputValue( + AdminQuotaSnowflakeHandler.QUOTA_INPUT_ID + ); + const newQuota = Number(quotaInput); + + if (isNaN(newQuota)) { + throw new BotException("Quota must be a number"); + } + if (newQuota < 0) { + throw new BotException("Quota must be positive"); + } + + await this.setQuota(interaction, snowflake, newQuota); + await reply(interaction, { + ephemeral: true, + embeds: [ + { + title: "Success", + description: "New quota set successfully!", + color: Colors.Green, + }, + ], + }); + } + + private handleError(interaction: RepliableInteraction, error: unknown) { + if ( + error instanceof DiscordjsError && + error.code === DiscordjsErrorCodes.InteractionCollectorError + ) { + setupInteractionCleanup(interaction, { time: 1 }); + } else if (error instanceof BotException) { + reply(interaction, { + ephemeral: true, + embeds: [ + { + title: "Error", + description: error.message, + color: Colors.Red, + }, + ], + components: [], + }); + } else { + throw error; + } + } + + protected abstract getCurrentQuota( + interaction: Interaction, + snowflake: string + ): Awaitable; + + protected abstract getInitialReply( + interaction: Interaction, + snowflake: string + ): Awaitable | InteractionResponse | null>; + + protected abstract deleteQuota( + interaction: Interaction, + snowflake: string + ): Awaitable; + + protected abstract getModalTitle(): string; + + protected abstract setQuota( + interaction: Interaction, + snowflake: string, + quota: number + ): Awaitable; +} diff --git a/packages/discord/src/interaction-handlers/admin/quota/AdminQuotaUserHandler.ts b/packages/discord/src/interaction-handlers/admin/quota/AdminQuotaUserHandler.ts index a090d90..e00cc6f 100644 --- a/packages/discord/src/interaction-handlers/admin/quota/AdminQuotaUserHandler.ts +++ b/packages/discord/src/interaction-handlers/admin/quota/AdminQuotaUserHandler.ts @@ -1,59 +1,56 @@ import { Interaction, Awaitable, - UserSelectMenuInteraction, - ModalSubmitInteraction, Colors, italic, ButtonBuilder, ButtonStyle, - DiscordjsError, - DiscordjsErrorCodes, - RepliableInteraction, - ModalBuilder, - ButtonInteraction, - TextInputStyle, + InteractionResponse, + Message, } from "discord.js"; -import InteractionHandler from "../../InteractionHandler.js"; -import isBotAdmin from "../../../utils/isBotAdmin.js"; import getHandlerId from "../../../utils/getHandlerId.js"; import reply from "../../../utils/reply.js"; -import setupInteractionCleanup from "../../../utils/setupInteractionCleanup.js"; -import { TextInputBuilder } from "@discordjs/builders"; -import BotException from "../../../exceptions/BotException.js"; -import { DEFAULT_INTERACTION_WAIT } from "../../../config/constants.js"; -import UserIdModalHandler from "../../UserIdModalHandler.js"; +import AdminQuotaSnowflakeHandler from "./AdminQuotaSnowflakeHandler.js"; -export default class AdminQuotaUserHandler extends InteractionHandler { +export default class AdminQuotaUserHandler extends AdminQuotaSnowflakeHandler { public static readonly ID = getHandlerId(AdminQuotaUserHandler.name); - private static readonly BUTTON_SET_ID = `${AdminQuotaUserHandler.ID}_button-set`; - private static readonly BUTTON_RESET_ID = `${AdminQuotaUserHandler.ID}_button-reset`; - private static readonly MODAL_SET_ID = `${AdminQuotaUserHandler.ID}_modal-set`; - private static readonly QUOTA_INPUT_ID = `${AdminQuotaUserHandler.ID}_quota-input`; + + constructor() { + super(AdminQuotaUserHandler.ID); + } public get name(): string { return AdminQuotaUserHandler.name; } - protected canHandle(interaction: Interaction): Awaitable { - return ( - isBotAdmin(interaction.user.id) && - (interaction.isUserSelectMenu() || interaction.isModalSubmit()) && - interaction.customId === AdminQuotaUserHandler.ID - ); + protected getCurrentQuota( + interaction: Interaction, + userId: string + ): Awaitable { + const { quotaManager } = interaction.client; + return quotaManager.getQuota(userId); } - protected async handle( - interaction: UserSelectMenuInteraction | ModalSubmitInteraction - ): Promise { + protected async getInitialReply( + interaction: Interaction, + userId: string + ): Promise | InteractionResponse | null> { const { quotaManager } = interaction.client; - const userId = interaction.isModalSubmit() - ? interaction.fields.getTextInputValue(UserIdModalHandler.INPUT_ID) - : interaction.values[0]; - const quota = await quotaManager.getQuota(userId); - const defaultQuota = await quotaManager.getDefaultQuota(); - const hasQuota = await quotaManager.hasQuota(userId); + const quota = await this.getCurrentQuota(interaction, userId); + const hasQuota = await quotaManager.hasUserQuota(userId); + const [quotaRole] = await quotaManager.getUserQuotaRole(userId); + + const quotaFormat = Intl.NumberFormat("en-US", { + style: "decimal", + maximumFractionDigits: 0, + }).format(quota); + const quotaType = (() => { + if (hasQuota) { + return ""; + } + return `(${italic(quotaRole?.name ?? "default")})`; + })(); const usage = await quotaManager.getUsage(userId); const usagePercent = new Intl.NumberFormat("en-US", { @@ -63,7 +60,7 @@ export default class AdminQuotaUserHandler extends InteractionHandler { const row = this.createMessageActionRow().addComponents( new ButtonBuilder() - .setCustomId(AdminQuotaUserHandler.BUTTON_SET_ID) + .setCustomId(AdminQuotaSnowflakeHandler.BUTTON_SET_ID) .setLabel("Set Quota") .setStyle(ButtonStyle.Primary) .setEmoji("🔧") @@ -72,24 +69,22 @@ export default class AdminQuotaUserHandler extends InteractionHandler { if (hasQuota) { row.addComponents( new ButtonBuilder() - .setCustomId(AdminQuotaUserHandler.BUTTON_RESET_ID) + .setCustomId(AdminQuotaSnowflakeHandler.BUTTON_RESET_ID) .setLabel("Reset Quota") .setStyle(ButtonStyle.Primary) .setEmoji("🔄") ); } - const response = await reply(interaction, { + return reply(interaction, { ephemeral: true, embeds: [ { - title: "Quota Information", + title: "User Quota Information", fields: [ { name: "Current Quota", - value: `${quota}${ - hasQuota ? "" : ` ${italic("(default)")}` - }`, + value: `${quotaFormat} ${quotaType}`, inline: true, }, { @@ -99,148 +94,30 @@ export default class AdminQuotaUserHandler extends InteractionHandler { }, ], color: Colors.Blue, - footer: { - text: `FYI: The default quota is ${defaultQuota} and can only be changed via the "DEFAULT_QUOTA" environment variable of the bot.`, - }, }, ], components: [row], }); - - if (!response) throw new Error("Failed to send message"); - - try { - const buttonInteraction = await response.awaitMessageComponent({ - filter: (i) => - i.user.id === interaction.user.id && - [ - AdminQuotaUserHandler.BUTTON_SET_ID, - AdminQuotaUserHandler.BUTTON_RESET_ID, - ].includes(i.customId), - time: DEFAULT_INTERACTION_WAIT, - }); - await response.delete(); - - if (!buttonInteraction.isButton()) - throw new Error("Expected button interaction"); - - await this.handleButtonInteraction(buttonInteraction, userId); - } catch (e) { - this.handleError(interaction, e); - } } - private async handleButtonInteraction( - interaction: ButtonInteraction, + protected async deleteQuota( + interaction: Interaction, userId: string - ) { + ): Promise { const { quotaManager } = interaction.client; - - switch (interaction.customId) { - case AdminQuotaUserHandler.BUTTON_SET_ID: - await this.showQuotaModal(interaction, userId); - break; - case AdminQuotaUserHandler.BUTTON_RESET_ID: - await quotaManager.deleteQuota(userId); - const defaultQuota = await quotaManager.getDefaultQuota(); - await reply(interaction, { - ephemeral: true, - embeds: [ - { - title: "Success", - description: `Quota was reset to the default value! (${defaultQuota})`, - color: Colors.Green, - }, - ], - }); - break; - default: - throw new Error( - `Unknown button custom id: ${interaction.customId}` - ); - } + await quotaManager.deleteUserQuota(userId); } - private async showQuotaModal( - interaction: ButtonInteraction, - userId: string - ) { - const { quotaManager } = interaction.client; - const quota = await quotaManager.getQuota(userId); - - const modal = new ModalBuilder() - .setCustomId(AdminQuotaUserHandler.MODAL_SET_ID) - .setTitle("Set User Quota") - .addComponents( - this.createModalActionRow().addComponents( - new TextInputBuilder() - .setCustomId(AdminQuotaUserHandler.QUOTA_INPUT_ID) - .setLabel("Quota") - .setPlaceholder(quota.toString()) - .setRequired(true) - .setMinLength(1) - .setMaxLength(10) - .setStyle(TextInputStyle.Short) - .setValue(quota.toString()) - ) - ); - - await interaction.showModal(modal); - - const modalInteraction = await interaction.awaitModalSubmit({ - filter: (i) => - i.user.id === interaction.user.id && - i.isModalSubmit() && - i.customId === AdminQuotaUserHandler.MODAL_SET_ID, - time: DEFAULT_INTERACTION_WAIT, - }); - await modalInteraction.deferUpdate(); - - const quotaInput = modalInteraction.fields.getTextInputValue( - AdminQuotaUserHandler.QUOTA_INPUT_ID - ); - const newQuota = Number(quotaInput); - - if (isNaN(newQuota)) { - throw new BotException("Quota must be a number"); - } - if (quota < 0) { - throw new BotException("Quota must be positive"); - } - - await quotaManager.setQuota(userId, newQuota); - await reply(interaction, { - ephemeral: true, - embeds: [ - { - title: "Success", - description: "New quota set successfully!", - color: Colors.Green, - }, - ], - }); + protected getModalTitle(): string { + return "Set User Quota"; } - private handleError(interaction: RepliableInteraction, error: unknown) { - if ( - error instanceof DiscordjsError && - error.code === DiscordjsErrorCodes.InteractionCollectorError - ) { - setupInteractionCleanup(interaction, { time: 1 }); - } else if (error instanceof BotException) { - reply(interaction, { - ephemeral: true, - embeds: [ - { - title: "Error", - description: error.message, - color: Colors.Red, - }, - ], - components: [], - }); - } else { - throw error; - } + protected async setQuota( + interaction: Interaction, + userId: string, + quota: number + ): Promise { + const { quotaManager } = interaction.client; + await quotaManager.setUserQuota(userId, quota); } } diff --git a/packages/discord/src/interaction-handlers/admin/usage-reset/AdminUsageResetUserHandler.ts b/packages/discord/src/interaction-handlers/admin/usage-reset/AdminUsageResetUserHandler.ts index 98179a1..81aea30 100644 --- a/packages/discord/src/interaction-handlers/admin/usage-reset/AdminUsageResetUserHandler.ts +++ b/packages/discord/src/interaction-handlers/admin/usage-reset/AdminUsageResetUserHandler.ts @@ -9,7 +9,7 @@ import InteractionHandler from "../../InteractionHandler.js"; import isBotAdmin from "../../../utils/isBotAdmin.js"; import getHandlerId from "../../../utils/getHandlerId.js"; import reply from "../../../utils/reply.js"; -import UserIdModalHandler from "../../UserIdModalHandler.js"; +import UserIdModalHandler from "../../SnowflakeModalHandler.js"; export default class AdminUsageResetUserHandler extends InteractionHandler { public static readonly ID = getHandlerId(AdminUsageResetUserHandler.name); diff --git a/packages/discord/src/managers/QuotaManager.ts b/packages/discord/src/managers/QuotaManager.ts index b8654cf..ed1c8a3 100644 --- a/packages/discord/src/managers/QuotaManager.ts +++ b/packages/discord/src/managers/QuotaManager.ts @@ -10,6 +10,8 @@ import { Conversation } from "gpt-turbo"; import { ConversationUser } from "../utils/types.js"; import isValidSnowflake from "../utils/isValidSnowflake.js"; import BotException from "../exceptions/BotException.js"; +import GPTTurboClient from "../GPTTurboClient.js"; +import { Role } from "discord.js"; export type DbType = "mongodb" | "mysql" | "postgres"; export default class QuotaManager< @@ -22,11 +24,18 @@ export default class QuotaManager< "__default__gptturbodiscord__quota__"; private dbType: TDbType = null as TDbType; - private quotas: TQuotas = null as TQuotas; + private userQuotas: TQuotas = null as TQuotas; + private roleQuotas: TQuotas = null as TQuotas; private usages: TUsages = null as TUsages; - constructor() { - this.quotas = this.getDb("gpt-turbo-discord-quotas") as TQuotas; + /** ID of each role that appear in the roleQuotas database */ + private quotedRoleIds: Set = new Set(); + /** ID of each guild that contain a quoted role from the quotedRoleIds */ + private quotedGuildIds: Set = new Set(); + + constructor(private client: GPTTurboClient) { + this.userQuotas = this.getDb("gpt-turbo-discord-userquotas") as TQuotas; + this.roleQuotas = this.getDb("gpt-turbo-discord-rolequotas") as TQuotas; this.usages = this.getDb("gpt-turbo-discord-usages") as TUsages; console.info( @@ -38,48 +47,148 @@ export default class QuotaManager< public async init() { if (!this.isEnabled()) return; - await this.quotas.set(QuotaManager.DEFAULT_QUOTA_KEY, DEFAULT_QUOTA); + await this.userQuotas.set( + QuotaManager.DEFAULT_QUOTA_KEY, + DEFAULT_QUOTA + ); } - public isEnabled(): this is QuotaManager { - return this.quotas !== null && this.usages !== null; + public async postInit() { + await this.refreshQuotedIds(); } - public async hasQuota(userId: string): Promise { - if (!this.isEnabled()) throw new BotException("Quotas are disabled"); - const id = this.parseUserId(userId); - return this.quotas.has(id); + public isEnabled(): this is QuotaManager { + return ( + this.userQuotas !== null && + this.roleQuotas !== null && + this.usages !== null + ); } public async getDefaultQuota(): Promise { if (!this.isEnabled()) throw new BotException("Quotas are disabled"); - const quota = await this.quotas.get(QuotaManager.DEFAULT_QUOTA_KEY); + const quota = await this.userQuotas.get(QuotaManager.DEFAULT_QUOTA_KEY); if (quota === undefined) throw new Error("No default quota found"); return quota; } - public async getQuota(userId: string): Promise { + /** + * Gets the quota assigned to a specific user in the database, if it exists, or null if it doesn't. + * **Important:** This does not take into account the user's role quotas. Use `getQuota` instead if you want to get the actual computed quota for a user. + */ + public async getUserQuota(userId: string): Promise { if (!this.isEnabled()) throw new BotException("Quotas are disabled"); - const id = this.parseUserId(userId); - return (await this.quotas.get(id)) ?? (await this.getDefaultQuota()); + const id = this.parseSnowflake(userId); + return (await this.userQuotas.get(id)) ?? null; } - public async setQuota(userId: string, quota: number) { + public async hasUserQuota(userId: string): Promise { + if (!this.isEnabled()) throw new BotException("Quotas are disabled"); + const id = this.parseSnowflake(userId); + return this.userQuotas.has(id); + } + + public async setUserQuota(userId: string, quota: number) { if (!this.isEnabled()) throw new BotException("Quotas are disabled"); if (quota < 0) throw new BotException("Quota must be positive"); - const id = this.parseUserId(userId); - await this.quotas.set(id, quota); + const id = this.parseSnowflake(userId); + await this.userQuotas.set(id, quota); + } + + public async deleteUserQuota(userId: string) { + if (!this.isEnabled()) throw new BotException("Quotas are disabled"); + const id = this.parseSnowflake(userId); + await this.userQuotas.delete(id); } - public async deleteQuota(userId: string) { + /** + * Gets the quota assigned to a role. + */ + public async getRoleQuota(roleId: string): Promise { if (!this.isEnabled()) throw new BotException("Quotas are disabled"); - const id = this.parseUserId(userId); - await this.quotas.delete(id); + const id = this.parseSnowflake(roleId); + return (await this.roleQuotas.get(id)) ?? null; + } + + public async hasRoleQuota(roleId: string): Promise { + if (!this.isEnabled()) throw new BotException("Quotas are disabled"); + const id = this.parseSnowflake(roleId); + return this.roleQuotas.has(id); + } + + /** + * If the user's quota is defined by a role quota, this method returns that role and the quota assigned to it. + */ + public async getUserQuotaRole( + userId: string + ): Promise<[Role | null, number]> { + let highestRole: Role | null = null; + let highestQuota = -1; + + const userQuota = await this.getUserQuota(userId); + if (userQuota !== null) return [highestRole, userQuota]; + + // find all guilds in quotedGuildIds that the user is in + const quotedGuilds = this.client.guilds.cache.filter((guild) => + this.quotedGuildIds.has(guild.id) + ); + const userGuilds = quotedGuilds.filter((guild) => + guild.members.cache.has(userId) + ); + + // find all roles in quotedRoleIds that the user has + const quotedRoles = userGuilds.flatMap((guild) => + guild.roles.cache.filter((role) => this.quotedRoleIds.has(role.id)) + ); + const userRoles = quotedRoles.filter((role) => + role.members.has(userId) + ); + + for (const role of userRoles.values()) { + const quota = await this.getRoleQuota(role.id); + if (quota !== null && quota > highestQuota) { + highestRole = role; + highestQuota = quota; + } + } + + return [highestRole, highestQuota]; + } + + public async setRoleQuota(roleId: string, quota: number) { + if (!this.isEnabled()) throw new BotException("Quotas are disabled"); + if (quota < 0) throw new BotException("Quota must be positive"); + const id = this.parseSnowflake(roleId); + await this.roleQuotas.set(id, quota); + await this.refreshQuotedIds(); + } + + public async deleteRoleQuota(roleId: string) { + if (!this.isEnabled()) throw new BotException("Quotas are disabled"); + const id = this.parseSnowflake(roleId); + await this.roleQuotas.delete(id); + await this.refreshQuotedIds(); + } + + /** + * This combines `getUserQuota` and `getRoleQuota` to determine the actual quota for a user. + * - User quotas take priority over role quotas + * - The highest role quota is used, if any + * - If no quota is found, the default quota is used + */ + public async getQuota(userId: string) { + const userQuota = await this.getUserQuota(userId); + if (userQuota !== null) return userQuota; + + const [role, roleQuota] = await this.getUserQuotaRole(userId); + if (role !== null) return roleQuota; + + return this.getDefaultQuota(); } public async getUsage(userId: string): Promise { if (!this.isEnabled()) throw new BotException("Quotas are disabled"); - const id = this.parseUserId(userId); + const id = this.parseSnowflake(userId); const usage = await this.usages.get(id); return usage ?? 0; } @@ -88,7 +197,7 @@ export default class QuotaManager< if (!this.isEnabled()) throw new BotException("Quotas are disabled"); if (usage < 0) throw new BotException("Usage must be positive"); const quota = await this.getQuota(userId); - const id = this.parseUserId(userId); + const id = this.parseSnowflake(userId); await this.usages.set(id, Math.min(usage, quota)); } @@ -104,7 +213,7 @@ export default class QuotaManager< const { user } = conversation.getConfig(); if (!user) throw new Error("No user found in conversation config"); - const userId = this.parseUserId(user); + const userId = this.parseSnowflake(user); const quota = await this.getQuota(userId); const usage = await this.getUsage(userId); @@ -116,17 +225,17 @@ export default class QuotaManager< const { user } = conversation.getConfig(); if (!user) throw new Error("No user found in conversation config"); - const userId = this.parseUserId(user); + const userId = this.parseSnowflake(user); const usage = await this.getUsage(userId); await this.usages.set(userId, usage + conversation.getSize()); } - private parseUserId(user: string) { - if (isValidSnowflake(user)) return user; - if (!this.isValidConversationUser(user)) - throw new BotException("Invalid user id"); - return user.split("-")[1]; + private parseSnowflake(snowflake: string) { + if (isValidSnowflake(snowflake)) return snowflake; + if (!this.isValidConversationUser(snowflake)) + throw new BotException("Invalid user/role id"); + return snowflake.split("-")[1]; } private isValidConversationUser(user: string): user is ConversationUser { @@ -176,4 +285,33 @@ export default class QuotaManager< return db; } + + private async refreshQuotedIds() { + if (!this.isEnabled()) return; + + await this.client.guilds.fetch(); + + const guilds = this.client.guilds.cache.map((guild) => guild); + const roles = guilds.flatMap((guild) => + guild.roles.cache.map((role) => role) + ); + + const tempQuotedRoleIds = new Set(); + const tempQuotedGuildIds = new Set(); + for await (const [roleId] of this.roleQuotas.iterator()) { + const quotedRole = roles.find((role) => role.id === roleId); + if (!quotedRole) { + console.warn( + "Role associated with a role quota not found:", + roleId + ); + continue; + } + tempQuotedRoleIds.add(roleId); + tempQuotedGuildIds.add(quotedRole.guild.id); + } + + this.quotedRoleIds = tempQuotedRoleIds; + this.quotedGuildIds = tempQuotedGuildIds; + } }