Refactor the main bot execution out of index.js
This commit is contained in:
parent
1c49e8b730
commit
cb304f522b
2 changed files with 150 additions and 118 deletions
148
src/execution.ts
Normal file
148
src/execution.ts
Normal file
|
|
@ -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<string, DiscordApi.Message[]> = 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<string, DiscordApi.Message> = 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);
|
||||
}
|
||||
120
src/index.ts
120
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<string, DiscordApi.Message[]> = 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<string, DiscordApi.Message> = 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);
|
||||
|
|
|
|||
Loading…
Reference in a new issue