Skip to content
This repository has been archived by the owner on Jun 2, 2024. It is now read-only.

Commit

Permalink
[discord] Role quotas (#12)
Browse files Browse the repository at this point in the history
  • Loading branch information
maxijonson committed Apr 29, 2023
1 parent 9bc4ed3 commit 7018d78
Show file tree
Hide file tree
Showing 13 changed files with 651 additions and 238 deletions.
4 changes: 2 additions & 2 deletions .github/PULL_REQUEST_TEMPLATE.md
Original file line number Diff line number Diff line change
Expand Up @@ -50,5 +50,5 @@ Issue: #

## Discord User (optional)

<!-- If your PR gets merged, we'll add this username to the "Contributors" role! Make sure to join the Discord beforehand so we can find you! -->
<!-- Don't forget the #! (e.g. MaxiJonson#1248) -->
<!-- If your PR gets merged, we'll add you to the "Contributors" role, which gives you access to a higher token quota to use on the server's bot! Make sure to join the Discord beforehand so we can find you! -->
<!-- Don't forget the # (e.g. MaxiJonson#1248) -->
5 changes: 3 additions & 2 deletions packages/discord/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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:

Expand All @@ -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

Expand Down
38 changes: 28 additions & 10 deletions packages/discord/src/GPTTurboClient.ts
Original file line number Diff line number Diff line change
Expand Up @@ -20,17 +20,18 @@ 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
> extends Client<Ready> {
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<string, Collection<string, number>>();
Expand All @@ -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"
)
);

Expand Down Expand Up @@ -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() {
Expand All @@ -117,4 +131,8 @@ export default class GPTTurboClient<
this.quotaManager.init(),
]);
}

private async postInit() {
await Promise.all([this.quotaManager.postInit()]);
}
}
2 changes: 1 addition & 1 deletion packages/discord/src/commands/admin.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down
34 changes: 27 additions & 7 deletions packages/discord/src/commands/usage.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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,
},
],
Expand Down
14 changes: 14 additions & 0 deletions packages/discord/src/components/RoleSelect.ts
Original file line number Diff line number Diff line change
@@ -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);
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -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<boolean> {
Expand All @@ -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)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,8 @@ import {
Interaction,
Awaitable,
StringSelectMenuInteraction,
ButtonBuilder,
ButtonStyle,
} from "discord.js";
import InteractionHandler from "../../InteractionHandler.js";
import isBotAdmin from "../../../utils/isBotAdmin.js";
Expand All @@ -12,21 +14,24 @@ 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<boolean> {
return (
isBotAdmin(interaction.user.id) &&
interaction.isStringSelectMenu() &&
interaction.customId === ADMIN_MENU_ID &&
interaction.values[0] === AdminQuotaUserMenuHandler.ID
interaction.values[0] === AdminQuotaMenuHandler.ID
);
}

Expand All @@ -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)
),
],
});
Expand Down
Original file line number Diff line number Diff line change
@@ -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<number> {
const { quotaManager } = interaction.client;
return (
(await quotaManager.getRoleQuota(roleId)) ??
(await quotaManager.getDefaultQuota())
);
}

protected async getInitialReply(
interaction: Interaction,
roleId: string
): Promise<Message<boolean> | InteractionResponse<boolean> | 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<void> {
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<void> {
const { quotaManager } = interaction.client;
await quotaManager.setRoleQuota(roleId, quota);
}
}
Loading

0 comments on commit 7018d78

Please sign in to comment.