diff --git a/src/execution.ts b/src/execution.ts deleted file mode 100644 index 3b4c3f1..0000000 --- a/src/execution.ts +++ /dev/null @@ -1,174 +0,0 @@ -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 the user that requested the execution - * @param request The request to get the user from - * @returns the user or guild member - */ -export function getAuthor(request: DiscordApi.Message | DiscordApi.Interaction) { - if (request instanceof DiscordApi.Message) return request.member ?? request.author; - 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 - */ -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 request the message to check and queue - */ -export async function queueRequest(request: DiscordApi.Message | DiscordApi.Interaction) { - if (!request.channelId) { - if (!(request instanceof DiscordApi.Message) && request.isRepliable()) - request.reply("request does not have channelId"); - console.log("There was incoming execution without channelId set, ignoring"); - console.log(request); - return; - } - - const userLimit = await getUserLimit(getAuthor(request), request.createdAt); - - if (userLimit !== false && userLimit.remaining <= 0) { - if (request instanceof DiscordApi.Message) { - request.react("🛑"); - if (!request.author.dmChannel) await request.author.createDM(); - request.author.dmChannel?.send({ - embeds: [{ - color: 0xff0000, - description: `You've used up your message limit for today, ${userLimit.limit} requests in last 24 hours`, - }] - }); - } - else if (request.isRepliable()) { - request.reply({ - content: `You've used up your message limit for today, ${userLimit.limit} requests in last 24 hours`, - ephemeral: true, - }); - } - return; - } - - const messagesForChannel = channelsRunning.ensure( - request.channelId, - () => { return []; }, - ); - const shouldStart = messagesForChannel.length == 0; - messagesForChannel.push(request); - if (shouldStart) - executeFromQueue(request.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 cba9b7b..27c214d 100644 --- a/src/index.ts +++ b/src/index.ts @@ -3,7 +3,8 @@ import { Configuration as OpenAIApiConfiguration, OpenAIApi } from "openai"; import { PrismaClient } from "@prisma/client"; import config from "./config"; -import { queueRequest } from "./execution"; +import toOpenAIMessages from "./toOpenAIMessages"; +import Moderation from "./moderation"; const discord = new DiscordApi.Client({ intents: [ @@ -23,11 +24,126 @@ 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; - queueRequest(message); + 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); }); +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);