diff --git a/src/execution.ts b/src/execution.ts new file mode 100644 index 0000000..13cffa4 --- /dev/null +++ b/src/execution.ts @@ -0,0 +1,148 @@ +import DiscordApi from "discord.js"; + +import { database, openai } from "./index"; +import Moderation from "./moderation"; +import config from "./config"; +import toOpenAIMessages from "./toOpenAIMessages"; + +/** Stores the queue requests on the channels. */ +const channelsRunning: DiscordApi.Collection = new DiscordApi.Collection(); + +/** + * 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 + */ +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: 25, remaining: 25 - usedLimit}; + + return {limit: userLimits.limit, remaining: userLimits.limit - usedLimit}; +} + +/** + * Check and queues up the request and runs it if there is nothing in queue. + * @param message the message to check and queue + */ +export async function queueRequest(message: DiscordApi.Message) { + + const userLimit = await getUserLimit(message.author, message.createdAt); + + if (userLimit !== false && userLimit.remaining <= 0) { + message.react("🛑"); + message.author.dmChannel?.send({ + embeds: [{ + color: 0xff0000, + description: `You've used up your message limit for today, ${userLimit.limit} requrests in last 24 hours` + }] + }); + return; + } + + const messagesForChannel = channelsRunning.ensure( + message.channelId, + () => { return [] as DiscordApi.Message[]; }, + ); + const shouldStart = messagesForChannel.length == 0; + messagesForChannel.push(message); + if (shouldStart) + executeFromQueue(message.channelId); +} + +/** + * Executes the queue for the channel + * @param channel the channel to run the queue for + */ +async function executeFromQueue(channel: string) { + const channelQueue = channelsRunning.get(channel) as DiscordApi.Message[]; + const message = channelQueue.at(0) as DiscordApi.Message; + + try { + let messages: DiscordApi.Collection = await message.channel.messages.fetch({ limit: config.limits.messages, cache: false }); + + messages = messages.filter(m => message.createdTimestamp - m.createdTimestamp < config.limits.time ); + + messages.forEach(m => Moderation.checkMessage(m)); + + message.channel.sendTyping(); + const answer = await openai.createChatCompletion({ + ...config.chatCompletionConfig, + messages: toOpenAIMessages(messages), + }); + + const usage = answer.data.usage; + if (usage != undefined) { + const channelName: string = message.inGuild() ? `${message.channel.name} (${message.guild.name})` : `@${message.author.tag}`; + console.log(`Used ${usage.total_tokens} (${usage.prompt_tokens} + ${usage.completion_tokens}) tokens for ${message.author.tag} (${message.author.id}) in #${channelName}`); + + database.usage.create({ + data: { + timestamp: message.createdAt, + user: BigInt(message.author.id), + channel: BigInt(message.channelId), + guild: message.guildId ? BigInt(message.guildId) : null, + usageReguest: usage.prompt_tokens, + usageResponse: usage.completion_tokens + } + }).catch((e => { + console.error("Failed to push to a database"); + console.error(e); + })); + } + + const answerContent = answer.data.choices[0].message?.content; + + if (answerContent != undefined && answerContent != "") { + const response = message.reply({ + content: answerContent, + allowedMentions: { + repliedUser: false, + } + }); + + Moderation.checkMessage(await response); + } + else { + message.react("😶"); + } + } catch (e) { + console.error(`Error ocurred while handling chat completion request (${(e as object).constructor.name}):`); + console.error(e); + + message.reply({ + embeds: [{ + color: 0xff0000, + description: "Something bad happened! :frowning:" + }], + allowedMentions: { + repliedUser: false, + } + }); + } + + channelQueue.shift(); + if (channelQueue.length == 0) + channelsRunning.delete(channel); + else + executeFromQueue(channel); +} diff --git a/src/index.ts b/src/index.ts index 27c214d..cba9b7b 100644 --- a/src/index.ts +++ b/src/index.ts @@ -3,8 +3,7 @@ import { Configuration as OpenAIApiConfiguration, OpenAIApi } from "openai"; import { PrismaClient } from "@prisma/client"; import config from "./config"; -import toOpenAIMessages from "./toOpenAIMessages"; -import Moderation from "./moderation"; +import { queueRequest } from "./execution"; const discord = new DiscordApi.Client({ intents: [ @@ -24,126 +23,11 @@ discord.on("ready", async event => { console.log(`Connected to Discord as ${event.user.tag} (${event.user.id})`); }); -const channelsRunning: DiscordApi.Collection = new DiscordApi.Collection(); - discord.on("messageCreate", async message => { if (message.author.bot) return; if (!message.mentions.has(message.client.user)) return; - const userLimits = await database.limits.findUnique({ - where: { - user: BigInt(message.author.id) - } - }); - - let limit = 25; - - if (userLimits?.limit) { - limit = userLimits.limit; - } - - if (!userLimits?.vip) { - const usedLimit = await database.usage.count({ - select: { _all: true }, - where: { - user: BigInt(message.author.id), - timestamp: { - gte: new Date(message.createdTimestamp - 1000*60*60*24 /*24 hours */) - } - }, - }); - - if (usedLimit._all >= limit) { - message.react("🛑"); - message.author.dmChannel?.send({ - embeds: [{ - color: 0xff0000, - description: `You've used up your message limit for today, ${limit} requrests in last 24 hours` - }] - }); - return; - } - } - - const messagesForChannel = channelsRunning.ensure(message.channelId, () => {return [] as DiscordApi.Message[];} ); - const shouldStart = messagesForChannel.length == 0; - messagesForChannel.push(message); - if (shouldStart) - onMessage(message.channelId); + queueRequest(message); }); -async function onMessage(channel: string) { - const channelQueue = channelsRunning.get(channel) as DiscordApi.Message[]; - const message = channelQueue.at(0) as DiscordApi.Message; - - try { - let messages: DiscordApi.Collection = await message.channel.messages.fetch({ limit: config.limits.messages, cache: false }); - - messages = messages.filter(m => message.createdTimestamp - m.createdTimestamp < config.limits.time ); - - messages.forEach(m => Moderation.checkMessage(m)); - - message.channel.sendTyping(); - const answer = await openai.createChatCompletion({ - ...config.chatCompletionConfig, - messages: toOpenAIMessages(messages), - }); - - const usage = answer.data.usage; - if (usage != undefined) { - const channelName: string = message.inGuild() ? `${message.channel.name} (${message.guild.name})` : `@${message.author.tag}`; - console.log(`Used ${usage.total_tokens} (${usage.prompt_tokens} + ${usage.completion_tokens}) tokens for ${message.author.tag} (${message.author.id}) in #${channelName}`); - - database.usage.create({ - data: { - timestamp: message.createdAt, - user: BigInt(message.author.id), - channel: BigInt(message.channelId), - guild: message.guildId ? BigInt(message.guildId) : null, - usageReguest: usage.prompt_tokens, - usageResponse: usage.completion_tokens - } - }).catch((e => { - console.error("Failed to push to a database"); - console.error(e); - })); - } - - const answerContent = answer.data.choices[0].message?.content; - - if (answerContent != undefined && answerContent != "") { - const response = message.reply({ - content: answerContent, - allowedMentions: { - repliedUser: false, - } - }); - - Moderation.checkMessage(await response); - } - else { - message.react("😶"); - } - } catch (e) { - console.error(`Error ocurred while handling chat completion request (${(e as object).constructor.name}):`); - console.error(e); - - message.reply({ - embeds: [{ - color: 0xff0000, - description: "Something bad happened! :frowning:" - }], - allowedMentions: { - repliedUser: false, - } - }); - } - - channelQueue.shift(); - if (channelQueue.length == 0) - channelsRunning.delete(channel); - else - onMessage(channel); -} - discord.login(config.tokens.Discord);